copy_wavelet_spec.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "print(sys.path)\n",
    "sys.path.append(\"C:/Users/jz421/Desktop/GlobalLocal/IEEG_Pipelines/\") #need to do this cuz otherwise ieeg isn't added to path...\n",
    "\n",
    "import mne.time_frequency\n",
    "import mne\n",
    "from ieeg.io import get_data, raw_from_layout\n",
    "from ieeg.navigate import trial_ieeg, channel_outlier_marker, crop_empty_data, outliers_to_nan\n",
    "from ieeg.calc.scaling import rescale\n",
    "import os\n",
    "from ieeg.timefreq.utils import wavelet_scaleogram, crop_pad\n",
    "import numpy as np\n",
    "\n",
    "from misc_functions import calculate_RTs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### get wavelets for incongruent and congruent\n",
    "this requires concatenating across trials for the events in the events list, note this is different from aaron's method below where each event is plotted separately instead of being concatenated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HOME = os.path.expanduser(\"~\")\n",
    "\n",
    "# get box directory depending on OS\n",
    "if os.name == 'nt': # windows\n",
    "    LAB_root = os.path.join(HOME, \"Box\", \"CoganLab\")\n",
    "else: # mac\n",
    "    LAB_root = os.path.join(HOME, \"Library\", \"CloudStorage\", \"Box-Box\", \"CoganLab\")\n",
    "\n",
    "layout = get_data(\"GlobalLocal\", root=LAB_root)\n",
    "\n",
    "output_name = \"Stimulus_i25and75_fixationCrossBase_1sec\"\n",
    "\n",
    "events = [\"Stimulus/c75\", \"Stimulus/c25\"]\n",
    "event_times = (-1,2)\n",
    "\n",
    "subjects = layout.get(return_type=\"id\", target=\"subject\")\n",
    "subjects = [\"D0057\", \"D0059\", \"D0063\"]\n",
    "\n",
    "for sub in subjects:\n",
    "\n",
    "    # Load the data\n",
    "    filt = raw_from_layout(layout.derivatives['derivatives/clean'], subject=sub,\n",
    "                            extension='.edf', desc='clean', preload=False) #get line-noise filtered data\n",
    "    print(filt)\n",
    "\n",
    "    ## Crop raw data to minimize processing time\n",
    "    good = crop_empty_data(filt)\n",
    "\n",
    "    # good.drop_channels(good.info['bads'])\n",
    "    good.info['bads'] = channel_outlier_marker(good, 3, 2)\n",
    "    good.drop_channels(good.info['bads'])\n",
    "    # good.info['bads'] += channel_outlier_marker(good, 4, 2)\n",
    "    # good.drop_channels(good.info['bads'])\n",
    "    good.load_data()\n",
    "\n",
    "    ch_type = filt.get_channel_types(only_data_chs=True)[0]\n",
    "    good.set_eeg_reference(ref_channels=\"average\", ch_type=ch_type)\n",
    "\n",
    "    # Remove intermediates from mem\n",
    "    good.plot()\n",
    "\n",
    "    ## epoching and trial outlier removal\n",
    "\n",
    "    save_dir = os.path.join(layout.root, 'derivatives', 'spec', 'wavelet', sub)\n",
    "    if not os.path.exists(save_dir):\n",
    "            os.makedirs(save_dir)\n",
    "\n",
    "    RTs, skipped = calculate_RTs(good)\n",
    "    avg_RT = np.median(RTs)\n",
    "    print(avg_RT)\n",
    "    # make stimulus baseline EpochsTFR\n",
    "    # times=[-1.5, 3.5] #this is for 0.5 sec of padding on each side. Just do 3 seconds as estimate of the full trial. Don't use avg RT because its different across subjects and makes later analysis hard.\n",
    "    times = [-1.5, 0.5] #padd 0.5 secs each side, want 1 seconds before stim onset\n",
    "    print(\"times: \" + str(times))\n",
    "    trials = trial_ieeg(good, \"Stimulus\", times, preload=True)\n",
    "\n",
    "    # print(\"base trials shape is: \" + str(trials._data.shape))\n",
    "    outliers_to_nan(trials, outliers=10)\n",
    "    base = wavelet_scaleogram(trials, n_jobs=1, decim=int(good.info['sfreq'] / 100))\n",
    "    crop_pad(base, \"0.5s\")\n",
    "\n",
    "\n",
    "    print(\"done with base\")\n",
    "\n",
    "    all_epochs_list = []\n",
    "\n",
    "    for event in events:\n",
    "    # Epoching and HG extraction for each specified event. Then concatenate all trials epochs objects together (do Stimulus/c25 and Stimulus/c75 for example, and combine to get all congruent trials)\n",
    "        times_adj = [times[0] - 0.5, times[1] + 0.5]\n",
    "        trials = trial_ieeg(good, event, times_adj, preload=True,\n",
    "                            reject_by_annotation=False)\n",
    "        all_epochs_list.append(trials)\n",
    "\n",
    "    # Concatenate all trials\n",
    "    all_trials = mne.concatenate_epochs(all_epochs_list)\n",
    "\n",
    "    times = [None, None]\n",
    "    times[0] = event_times[0] - 0.5\n",
    "    times[1] = event_times[1] + 0.5\n",
    "\n",
    "    outliers_to_nan(all_trials, outliers=10)\n",
    "    spec = wavelet_scaleogram(all_trials, n_jobs=1, decim=int(good.info['sfreq'] / 100))\n",
    "    crop_pad(spec, \"0.5s\")\n",
    "    spec_a = rescale(spec, base, copy=True, mode='ratio').average(\n",
    "            lambda x: np.nanmean(x, axis=0), copy=True)\n",
    "    spec_a._data = np.log10(spec_a._data) * 20\n",
    "    fnames = [os.path.relpath(f, layout.root) for f in good.filenames]\n",
    "    spec_a.info['subject_info']['files'] = tuple(fnames)\n",
    "    spec_a.info['bads'] = good.info['bads']\n",
    "    filename = os.path.join(save_dir, f'{output_name}-tfr.h5')\n",
    "    mne.time_frequency.write_tfrs(filename, spec_a, overwrite=True)\n",
    "    spec_a.save(os.path.join(save_dir, f'{output_name}-avg.fif'), overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### this is for combined congruent or incongruent, it combines the events.  \n",
    "it clears the outputs after the first cell cuz no memory and then runs the rest of the subjects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "print(sys.path)\n",
    "sys.path.append(\"C:/Users/jz421/Desktop/GlobalLocal/IEEG_Pipelines/\") #need to do this cuz otherwise ieeg isn't added to path...\n",
    "\n",
    "import mne.time_frequency\n",
    "import mne\n",
    "from ieeg.io import get_data, raw_from_layout\n",
    "from ieeg.navigate import trial_ieeg, channel_outlier_marker, crop_empty_data, outliers_to_nan\n",
    "from ieeg.calc.scaling import rescale\n",
    "import os\n",
    "from ieeg.timefreq.utils import wavelet_scaleogram, crop_pad\n",
    "import numpy as np\n",
    "\n",
    "from misc_functions import calculate_RTs\n",
    "\n",
    "HOME = os.path.expanduser(\"~\")\n",
    "\n",
    "# get box directory depending on OS\n",
    "if os.name == 'nt': # windows\n",
    "    LAB_root = os.path.join(HOME, \"Box\", \"CoganLab\")\n",
    "else: # mac\n",
    "    LAB_root = os.path.join(HOME, \"Library\", \"CloudStorage\", \"Box-Box\", \"CoganLab\")\n",
    "\n",
    "layout = get_data(\"GlobalLocal\", root=LAB_root)\n",
    "\n",
    "output_name = \"Stimulus_i25and75_fixationCrossBase_1sec\"\n",
    "\n",
    "events = [\"Stimulus/i75\", \"Stimulus/i25\"]\n",
    "event_times = (-1,2)\n",
    "\n",
    "subjects = layout.get(return_type=\"id\", target=\"subject\")\n",
    "subjects = [\"D0065\", \"D0069\", \"D0071\"]\n",
    "\n",
    "for sub in subjects:\n",
    "\n",
    "    # Load the data\n",
    "    filt = raw_from_layout(layout.derivatives['derivatives/clean'], subject=sub,\n",
    "                            extension='.edf', desc='clean', preload=False) #get line-noise filtered data\n",
    "    print(filt)\n",
    "\n",
    "    ## Crop raw data to minimize processing time\n",
    "    good = crop_empty_data(filt)\n",
    "\n",
    "    # good.drop_channels(good.info['bads'])\n",
    "    good.info['bads'] = channel_outlier_marker(good, 3, 2)\n",
    "    good.drop_channels(good.info['bads'])\n",
    "    # good.info['bads'] += channel_outlier_marker(good, 4, 2)\n",
    "    # good.drop_channels(good.info['bads'])\n",
    "    good.load_data()\n",
    "\n",
    "    ch_type = filt.get_channel_types(only_data_chs=True)[0]\n",
    "    good.set_eeg_reference(ref_channels=\"average\", ch_type=ch_type)\n",
    "\n",
    "    # Remove intermediates from mem\n",
    "    good.plot()\n",
    "\n",
    "    ## epoching and trial outlier removal\n",
    "\n",
    "    save_dir = os.path.join(layout.root, 'derivatives', 'spec', 'wavelet', sub)\n",
    "    if not os.path.exists(save_dir):\n",
    "            os.makedirs(save_dir)\n",
    "\n",
    "    RTs, skipped = calculate_RTs(good)\n",
    "    avg_RT = np.median(RTs)\n",
    "    print(avg_RT)\n",
    "    # make stimulus baseline EpochsTFR\n",
    "    # times=[-1.5, 3.5] #this is for 0.5 sec of padding on each side. Just do 3 seconds as estimate of the full trial. Don't use avg RT because its different across subjects and makes later analysis hard.\n",
    "    times = [-1.5, 0.5] #padd 0.5 secs each side, want 1 seconds before stim onset\n",
    "    print(\"times: \" + str(times))\n",
    "    trials = trial_ieeg(good, \"Stimulus\", times, preload=True)\n",
    "\n",
    "    # print(\"base trials shape is: \" + str(trials._data.shape))\n",
    "    outliers_to_nan(trials, outliers=10)\n",
    "    base = wavelet_scaleogram(trials, n_jobs=1, decim=int(good.info['sfreq'] / 100))\n",
    "    crop_pad(base, \"0.5s\")\n",
    "\n",
    "\n",
    "    print(\"done with base\")\n",
    "\n",
    "    all_epochs_list = []\n",
    "\n",
    "    for event in events:\n",
    "    # Epoching and HG extraction for each specified event. Then concatenate all trials epochs objects together (do Stimulus/c25 and Stimulus/c75 for example, and combine to get all congruent trials)\n",
    "        times_adj = [times[0] - 0.5, times[1] + 0.5]\n",
    "        trials = trial_ieeg(good, event, times_adj, preload=True,\n",
    "                            reject_by_annotation=False)\n",
    "        all_epochs_list.append(trials)\n",
    "\n",
    "    # Concatenate all trials\n",
    "    all_trials = mne.concatenate_epochs(all_epochs_list)\n",
    "\n",
    "    times = [None, None]\n",
    "    times[0] = event_times[0] - 0.5\n",
    "    times[1] = event_times[1] + 0.5\n",
    "\n",
    "    outliers_to_nan(all_trials, outliers=10)\n",
    "    spec = wavelet_scaleogram(all_trials, n_jobs=1, decim=int(good.info['sfreq'] / 100))\n",
    "    crop_pad(spec, \"0.5s\")\n",
    "    spec_a = rescale(spec, base, copy=True, mode='ratio').average(\n",
    "            lambda x: np.nanmean(x, axis=0), copy=True)\n",
    "    spec_a._data = np.log10(spec_a._data) * 20\n",
    "    fnames = [os.path.relpath(f, layout.root) for f in good.filenames]\n",
    "    spec_a.info['subject_info']['files'] = tuple(fnames)\n",
    "    spec_a.info['bads'] = good.info['bads']\n",
    "    filename = os.path.join(save_dir, f'{output_name}-tfr.h5')\n",
    "    mne.time_frequency.write_tfrs(filename, spec_a, overwrite=True)\n",
    "    spec_a.save(os.path.join(save_dir, f'{output_name}-avg.fif'), overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import clear_output\n",
    "clear_output(wait=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "print(sys.path)\n",
    "sys.path.append(\"C:/Users/jz421/Desktop/GlobalLocal/IEEG_Pipelines/\") #need to do this cuz otherwise ieeg isn't added to path...\n",
    "\n",
    "import mne.time_frequency\n",
    "import mne\n",
    "from ieeg.io import get_data, raw_from_layout\n",
    "from ieeg.navigate import trial_ieeg, channel_outlier_marker, crop_empty_data, outliers_to_nan\n",
    "from ieeg.calc.scaling import rescale\n",
    "import os\n",
    "from ieeg.timefreq.utils import wavelet_scaleogram, crop_pad\n",
    "import numpy as np\n",
    "\n",
    "from misc_functions import calculate_RTs\n",
    "\n",
    "HOME = os.path.expanduser(\"~\")\n",
    "\n",
    "# get box directory depending on OS\n",
    "if os.name == 'nt': # windows\n",
    "    LAB_root = os.path.join(HOME, \"Box\", \"CoganLab\")\n",
    "else: # mac\n",
    "    LAB_root = os.path.join(HOME, \"Library\", \"CloudStorage\", \"Box-Box\", \"CoganLab\")\n",
    "\n",
    "layout = get_data(\"GlobalLocal\", root=LAB_root)\n",
    "\n",
    "output_name = \"Stimulus_i25and75_fixationCrossBase_1sec\"\n",
    "\n",
    "events = [\"Stimulus/i75\", \"Stimulus/i25\"]\n",
    "event_times = (-1,2)\n",
    "\n",
    "subjects = layout.get(return_type=\"id\", target=\"subject\")\n",
    "subjects = [\"D0057\", \"D0059\", \"D0063\"]\n",
    "\n",
    "for sub in subjects:\n",
    "\n",
    "    # Load the data\n",
    "    filt = raw_from_layout(layout.derivatives['derivatives/clean'], subject=sub,\n",
    "                            extension='.edf', desc='clean', preload=False) #get line-noise filtered data\n",
    "    print(filt)\n",
    "\n",
    "    ## Crop raw data to minimize processing time\n",
    "    good = crop_empty_data(filt)\n",
    "\n",
    "    # good.drop_channels(good.info['bads'])\n",
    "    good.info['bads'] = channel_outlier_marker(good, 3, 2)\n",
    "    good.drop_channels(good.info['bads'])\n",
    "    # good.info['bads'] += channel_outlier_marker(good, 4, 2)\n",
    "    # good.drop_channels(good.info['bads'])\n",
    "    good.load_data()\n",
    "\n",
    "    ch_type = filt.get_channel_types(only_data_chs=True)[0]\n",
    "    good.set_eeg_reference(ref_channels=\"average\", ch_type=ch_type)\n",
    "\n",
    "    # Remove intermediates from mem\n",
    "    good.plot()\n",
    "\n",
    "    ## epoching and trial outlier removal\n",
    "\n",
    "    save_dir = os.path.join(layout.root, 'derivatives', 'spec', 'wavelet', sub)\n",
    "    if not os.path.exists(save_dir):\n",
    "            os.makedirs(save_dir)\n",
    "\n",
    "    RTs, skipped = calculate_RTs(good)\n",
    "    avg_RT = np.median(RTs)\n",
    "    print(avg_RT)\n",
    "    # make stimulus baseline EpochsTFR\n",
    "    # times=[-1.5, 3.5] #this is for 0.5 sec of padding on each side. Just do 3 seconds as estimate of the full trial. Don't use avg RT because its different across subjects and makes later analysis hard.\n",
    "    times = [-1.5, 0.5] #padd 0.5 secs each side, want 1 seconds before stim onset\n",
    "    print(\"times: \" + str(times))\n",
    "    trials = trial_ieeg(good, \"Stimulus\", times, preload=True)\n",
    "\n",
    "    # print(\"base trials shape is: \" + str(trials._data.shape))\n",
    "    outliers_to_nan(trials, outliers=10)\n",
    "    base = wavelet_scaleogram(trials, n_jobs=1, decim=int(good.info['sfreq'] / 100))\n",
    "    crop_pad(base, \"0.5s\")\n",
    "\n",
    "\n",
    "    print(\"done with base\")\n",
    "\n",
    "    all_epochs_list = []\n",
    "\n",
    "    for event in events:\n",
    "    # Epoching and HG extraction for each specified event. Then concatenate all trials epochs objects together (do Stimulus/c25 and Stimulus/c75 for example, and combine to get all congruent trials)\n",
    "        times_adj = [times[0] - 0.5, times[1] + 0.5]\n",
    "        trials = trial_ieeg(good, event, times_adj, preload=True,\n",
    "                            reject_by_annotation=False)\n",
    "        all_epochs_list.append(trials)\n",
    "\n",
    "    # Concatenate all trials\n",
    "    all_trials = mne.concatenate_epochs(all_epochs_list)\n",
    "\n",
    "    times = [None, None]\n",
    "    times[0] = event_times[0] - 0.5\n",
    "    times[1] = event_times[1] + 0.5\n",
    "\n",
    "    outliers_to_nan(all_trials, outliers=10)\n",
    "    spec = wavelet_scaleogram(all_trials, n_jobs=1, decim=int(good.info['sfreq'] / 100))\n",
    "    crop_pad(spec, \"0.5s\")\n",
    "    spec_a = rescale(spec, base, copy=True, mode='ratio').average(\n",
    "            lambda x: np.nanmean(x, axis=0), copy=True)\n",
    "    spec_a._data = np.log10(spec_a._data) * 20\n",
    "    fnames = [os.path.relpath(f, layout.root) for f in good.filenames]\n",
    "    spec_a.info['subject_info']['files'] = tuple(fnames)\n",
    "    spec_a.info['bads'] = good.info['bads']\n",
    "    filename = os.path.join(save_dir, f'{output_name}-tfr.h5')\n",
    "    mne.time_frequency.write_tfrs(filename, spec_a, overwrite=True)\n",
    "    spec_a.save(os.path.join(save_dir, f'{output_name}-avg.fif'), overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2/22 get wavelets for stimulus using new baseline of 0.2 sec before stim onset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HOME = os.path.expanduser(\"~\")\n",
    "\n",
    "# get box directory depending on OS\n",
    "if os.name == 'nt': # windows\n",
    "    LAB_root = os.path.join(HOME, \"Box\", \"CoganLab\")\n",
    "else: # mac\n",
    "    LAB_root = os.path.join(HOME, \"Library\", \"CloudStorage\", \"Box-Box\", \"CoganLab\")\n",
    "\n",
    "layout = get_data(\"GlobalLocal\", root=LAB_root)\n",
    "\n",
    "\n",
    "layout = get_data(\"GlobalLocal\", root=LAB_root)\n",
    "subjects = layout.get(return_type=\"id\", target=\"subject\")\n",
    "\n",
    "print(subjects)\n",
    "\n",
    "subjects = ['D0063', 'D0065', 'D0069', 'D0071', 'D0077', 'D0090', 'D0094', 'D0100', 'D0102', 'D0103']\n",
    "for sub in subjects:\n",
    "        # Load the data\n",
    "        filt = raw_from_layout(layout.derivatives['derivatives/clean'], subject=sub,\n",
    "                                extension='.edf', desc='clean', preload=False) #get line-noise filtered data\n",
    "        print(filt)\n",
    "\n",
    "        ## Crop raw data to minimize processing time\n",
    "        good = crop_empty_data(filt)\n",
    "\n",
    "        # good.drop_channels(good.info['bads'])\n",
    "        good.info['bads'] = channel_outlier_marker(good, 3, 2)\n",
    "        good.drop_channels(good.info['bads'])\n",
    "        # good.info['bads'] += channel_outlier_marker(good, 4, 2)\n",
    "        # good.drop_channels(good.info['bads'])\n",
    "        good.load_data()\n",
    "\n",
    "        ch_type = filt.get_channel_types(only_data_chs=True)[0]\n",
    "        good.set_eeg_reference(ref_channels=\"average\", ch_type=ch_type)\n",
    "\n",
    "        # Remove intermediates from mem\n",
    "        good.plot()\n",
    "\n",
    "        ## epoching and trial outlier removal\n",
    "\n",
    "        save_dir = os.path.join(layout.root, 'derivatives', 'spec', 'wavelet', sub)\n",
    "        if not os.path.exists(save_dir):\n",
    "                os.makedirs(save_dir)\n",
    "\n",
    "        RTs, skipped = calculate_RTs(good)\n",
    "        avg_RT = np.median(RTs)\n",
    "\n",
    "        # make stimulus baseline EpochsTFR\n",
    "        times = [-0.7, 0.5] #padd 0.5 secs each side, want 0.2 seconds before stim onset\n",
    "        print(\"times: \" + str(times))\n",
    "        trials = trial_ieeg(good, \"Stimulus\", times, preload=True)\n",
    "\n",
    "        # print(\"base trials shape is: \" + str(trials._data.shape))\n",
    "        outliers_to_nan(trials, outliers=10)\n",
    "        base = wavelet_scaleogram(trials, n_jobs=1, decim=int(good.info['sfreq'] / 100))\n",
    "        crop_pad(base, \"0.5s\")\n",
    "\n",
    "\n",
    "        #now do rescale with the concatenated baseline epochs\n",
    "        for event, t in [(\"Stimulus\",(-1, 2))]:\n",
    "                times = [None, None]\n",
    "                times[0] = t[0] - 0.5\n",
    "                times[1] = t[1] + 0.5\n",
    "                trials = trial_ieeg(good, event, times, preload=True)\n",
    "                outliers_to_nan(trials, outliers=10)\n",
    "                spec = wavelet_scaleogram(trials, n_jobs=1, decim=int(good.info['sfreq'] / 100))\n",
    "                crop_pad(spec, \"0.5s\")\n",
    "                spec_a = rescale(spec, base, copy=True, mode='ratio').average(\n",
    "                        lambda x: np.nanmean(x, axis=0), copy=True)\n",
    "                spec_a._data = np.log10(spec_a._data) * 20\n",
    "                fnames = [os.path.relpath(f, layout.root) for f in good.filenames]\n",
    "                spec_a.info['subject_info']['files'] = tuple(fnames)\n",
    "                spec_a.info['bads'] = good.info['bads']\n",
    "                filename = os.path.join(save_dir, f'{event}_fixationCrossBase_0.2sec-tfr.h5')\n",
    "                mne.time_frequency.write_tfrs(filename, spec_a, overwrite=True)\n",
    "                spec_a.save(os.path.join(save_dir, f'{event}_fixationCrossBase_0.2sec-avg.fif'), overwrite=True)\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7/5 try to get wavelets for all stimulus/response  \n",
    "*note this is kinda old code, don't use it except for stimulus and response specifically"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HOME = os.path.expanduser(\"~\")\n",
    "\n",
    "# get box directory depending on OS\n",
    "if os.name == 'nt': # windows\n",
    "    LAB_root = os.path.join(HOME, \"Box\", \"CoganLab\")\n",
    "else: # mac\n",
    "    LAB_root = os.path.join(HOME, \"Library\", \"CloudStorage\", \"Box-Box\", \"CoganLab\")\n",
    "\n",
    "layout = get_data(\"GlobalLocal\", root=LAB_root)\n",
    "\n",
    "\n",
    "layout = get_data(\"GlobalLocal\", root=LAB_root)\n",
    "subjects = layout.get(return_type=\"id\", target=\"subject\")\n",
    "\n",
    "print(subjects)\n",
    "# subjects = ['D0077', 'D0090']\n",
    "for sub in subjects:\n",
    "        # Load the data\n",
    "        filt = raw_from_layout(layout.derivatives['derivatives/clean'], subject=sub,\n",
    "                                extension='.edf', desc='clean', preload=False) #get line-noise filtered data\n",
    "        print(filt)\n",
    "\n",
    "        ## Crop raw data to minimize processing time\n",
    "        good = crop_empty_data(filt)\n",
    "\n",
    "        # good.drop_channels(good.info['bads'])\n",
    "        good.info['bads'] = channel_outlier_marker(good, 3, 2)\n",
    "        good.drop_channels(good.info['bads'])\n",
    "        # good.info['bads'] += channel_outlier_marker(good, 4, 2)\n",
    "        # good.drop_channels(good.info['bads'])\n",
    "        good.load_data()\n",
    "\n",
    "        ch_type = filt.get_channel_types(only_data_chs=True)[0]\n",
    "        good.set_eeg_reference(ref_channels=\"average\", ch_type=ch_type)\n",
    "\n",
    "        # Remove intermediates from mem\n",
    "        good.plot()\n",
    "\n",
    "        ## epoching and trial outlier removal\n",
    "\n",
    "        save_dir = os.path.join(layout.root, 'derivatives', 'spec', 'wavelet', sub)\n",
    "        if not os.path.exists(save_dir):\n",
    "                os.makedirs(save_dir)\n",
    "\n",
    "        RTs, skipped = calculate_RTs(good)\n",
    "        avg_RT = np.median(RTs)\n",
    "        \n",
    "        # make stimulus baseline EpochsTFR\n",
    "        times = [-1.5, 0.5] #padd 0.5 secs each side, want 1 seconds before stim onset\n",
    "        print(\"times: \" + str(times))\n",
    "        trials = trial_ieeg(good, \"Stimulus\", times, preload=True)\n",
    "\n",
    "        # print(\"base trials shape is: \" + str(trials._data.shape))\n",
    "        outliers_to_nan(trials, outliers=10)\n",
    "        base = wavelet_scaleogram(trials, n_jobs=1, decim=int(good.info['sfreq'] / 100))\n",
    "        crop_pad(base, \"0.5s\")\n",
    "\n",
    "\n",
    "        #now do rescale with the concatenated baseline epochs\n",
    "        for event, t in zip((\"Stimulus\", \"Response\"),((-1, 2), (-1, 2))):\n",
    "                times = [None, None]\n",
    "                times[0] = t[0] - 0.5\n",
    "                times[1] = t[1] + 0.5\n",
    "                trials = trial_ieeg(good, event, times, preload=True)\n",
    "                outliers_to_nan(trials, outliers=10)\n",
    "                spec = wavelet_scaleogram(trials, n_jobs=1, decim=int(good.info['sfreq'] / 100))\n",
    "                crop_pad(spec, \"0.5s\")\n",
    "                spec_a = rescale(spec, base, copy=True, mode='ratio').average(\n",
    "                        lambda x: np.nanmean(x, axis=0), copy=True)\n",
    "                spec_a._data = np.log10(spec_a._data) * 20\n",
    "                fnames = [os.path.relpath(f, layout.root) for f in good.filenames]\n",
    "                spec_a.info['subject_info']['files'] = tuple(fnames)\n",
    "                spec_a.info['bads'] = good.info['bads']\n",
    "                filename = os.path.join(save_dir, f'{event}_fullTrialBase-tfr.h5')\n",
    "                mne.time_frequency.write_tfrs(filename, spec_a, overwrite=True)\n",
    "                spec_a.save(os.path.join(save_dir, f'{event}_fullTrialBase-avg.fif'), overwrite=True)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### testing with one subject"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HOME = os.path.expanduser(\"~\")\n",
    "\n",
    "# get box directory depending on OS\n",
    "if os.name == 'nt': # windows\n",
    "    LAB_root = os.path.join(HOME, \"Box\", \"CoganLab\")\n",
    "else: # mac\n",
    "    LAB_root = os.path.join(HOME, \"Library\", \"CloudStorage\", \"Box-Box\", \"CoganLab\")\n",
    "\n",
    "layout = get_data(\"GlobalLocal\", root=LAB_root)\n",
    "\n",
    "\n",
    "layout = get_data(\"GlobalLocal\", root=LAB_root)\n",
    "subjects = layout.get(return_type=\"id\", target=\"subject\")\n",
    "\n",
    "# subjects = [\"D0059\", \"D0063\", \"D0065\", \"D0069\", \"D0071\"]\n",
    "\n",
    "# for sub in subjects:\n",
    "\n",
    "sub='D0057' #why does this not work in the for loop???\n",
    "# Load the data\n",
    "filt = raw_from_layout(layout.derivatives['derivatives/clean'], subject=sub,\n",
    "                        extension='.edf', desc='clean', preload=False) #get line-noise filtered data\n",
    "print(filt)\n",
    "\n",
    "## Crop raw data to minimize processing time\n",
    "good = crop_empty_data(filt)\n",
    "\n",
    "# good.drop_channels(good.info['bads'])\n",
    "good.info['bads'] = channel_outlier_marker(good, 3, 2)\n",
    "good.drop_channels(good.info['bads'])\n",
    "# good.info['bads'] += channel_outlier_marker(good, 4, 2)\n",
    "# good.drop_channels(good.info['bads'])\n",
    "good.load_data()\n",
    "\n",
    "ch_type = filt.get_channel_types(only_data_chs=True)[0]\n",
    "good.set_eeg_reference(ref_channels=\"average\", ch_type=ch_type)\n",
    "\n",
    "# Remove intermediates from mem\n",
    "good.plot()\n",
    "\n",
    "## epoching and trial outlier removal\n",
    "\n",
    "save_dir = os.path.join(layout.root, 'derivatives', 'spec', 'wavelet', sub)\n",
    "if not os.path.exists(save_dir):\n",
    "        os.makedirs(save_dir)\n",
    "\n",
    "RTs, skipped = calculate_RTs(good)\n",
    "avg_RT = np.median(RTs)\n",
    "print(avg_RT)\n",
    "# make stimulus baseline EpochsTFR\n",
    "# times=[-1.5, 3.5] #this is for 0.5 sec of padding on each side. Just do 3 seconds as estimate of the full trial. Don't use avg RT because its different across subjects and makes later analysis hard.\n",
    "times = [-1.5, 0.5] #padd 0.5 secs each side, want 1 seconds before stim onset\n",
    "print(\"times: \" + str(times))\n",
    "trials = trial_ieeg(good, \"Stimulus\", times, preload=True)\n",
    "\n",
    "# print(\"base trials shape is: \" + str(trials._data.shape))\n",
    "outliers_to_nan(trials, outliers=10)\n",
    "base = wavelet_scaleogram(trials, n_jobs=1, decim=int(good.info['sfreq'] / 100))\n",
    "crop_pad(base, \"0.5s\")\n",
    "\n",
    "\n",
    "print(\"done with base\")\n",
    "\n",
    "#now do rescale with the concatenated baseline epochs\n",
    "for (event, t), filename_event in zip(zip((\"Stimulus/i25\", \"Stimulus/c25\"),((-1, 2), (-1, 2))), (\"Stimulus_i25\", \"Stimulus_c25\")):\n",
    "    times = [None, None]\n",
    "    times[0] = t[0] - 0.5\n",
    "    times[1] = t[1] + 0.5\n",
    "    trials = trial_ieeg(good, event, times, preload=True)\n",
    "    print(f\"trials for {event} is: \" + str(trials._data.shape))\n",
    "    outliers_to_nan(trials, outliers=10)\n",
    "    spec = wavelet_scaleogram(trials, n_jobs=1, decim=int(good.info['sfreq'] / 100))\n",
    "    crop_pad(spec, \"0.5s\")\n",
    "    spec_a = rescale(spec, base, copy=True, mode='ratio').average(\n",
    "            lambda x: np.nanmean(x, axis=0), copy=True)\n",
    "    spec_a._data = np.log10(spec_a._data) * 20\n",
    "    fnames = [os.path.relpath(f, layout.root) for f in good.filenames]\n",
    "    spec_a.info['subject_info']['files'] = tuple(fnames)\n",
    "    spec_a.info['bads'] = good.info['bads']\n",
    "    filename = os.path.join(save_dir, f'{sub}_{filename_event}_1secBeforeStimOnset-tfr.h5')\n",
    "    mne.time_frequency.write_tfrs(filename, spec_a, overwrite=True)\n",
    "    spec_a.save(os.path.join(save_dir, f'{sub}_{filename_event}_1secBeforeStimOnset-avg.fif'), overwrite=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filename = os.path.join(save_dir, f'{sub}_Stimulusi25_1secBeforeStimOnset-tfr.h5')\n",
    "mne.time_frequency.write_tfrs(filename, spec_a, overwrite=True)\n",
    "spec_a.save(os.path.join(save_dir, f'{sub}_Stimulusi25_1secBeforeStimOnset-avg.fif'), overwrite=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HOME = os.path.expanduser(\"~\")\n",
    "\n",
    "# get box directory depending on OS\n",
    "if os.name == 'nt': # windows\n",
    "    LAB_root = os.path.join(HOME, \"Box\", \"CoganLab\")\n",
    "else: # mac\n",
    "    LAB_root = os.path.join(HOME, \"Library\", \"CloudStorage\", \"Box-Box\", \"CoganLab\")\n",
    "\n",
    "layout = get_data(\"GlobalLocal\", root=LAB_root)\n",
    "\n",
    "\n",
    "layout = get_data(\"GlobalLocal\", root=LAB_root)\n",
    "subjects = layout.get(return_type=\"id\", target=\"subject\")\n",
    "\n",
    "subjects = [\"D0057\", \"D0059\", \"D0063\", \"D0065\", \"D0069\", \"D0071\"]\n",
    "\n",
    "for sub in subjects:\n",
    "\n",
    "    # Load the data\n",
    "    filt = raw_from_layout(layout.derivatives['derivatives/clean'], subject=sub,\n",
    "                            extension='.edf', desc='clean', preload=False) #get line-noise filtered data\n",
    "    print(filt)\n",
    "\n",
    "    ## Crop raw data to minimize processing time\n",
    "    good = crop_empty_data(filt)\n",
    "\n",
    "    # good.drop_channels(good.info['bads'])\n",
    "    good.info['bads'] = channel_outlier_marker(good, 3, 2)\n",
    "    good.drop_channels(good.info['bads'])\n",
    "    # good.info['bads'] += channel_outlier_marker(good, 4, 2)\n",
    "    # good.drop_channels(good.info['bads'])\n",
    "    good.load_data()\n",
    "\n",
    "    ch_type = filt.get_channel_types(only_data_chs=True)[0]\n",
    "    good.set_eeg_reference(ref_channels=\"average\", ch_type=ch_type)\n",
    "\n",
    "    # Remove intermediates from mem\n",
    "    good.plot()\n",
    "\n",
    "    ## epoching and trial outlier removal\n",
    "\n",
    "    save_dir = os.path.join(layout.root, 'derivatives', 'spec', 'wavelet', sub)\n",
    "    if not os.path.exists(save_dir):\n",
    "            os.makedirs(save_dir)\n",
    "\n",
    "    RTs, skipped = calculate_RTs(good)\n",
    "    avg_RT = np.median(RTs)\n",
    "    print(avg_RT)\n",
    "    # make stimulus baseline EpochsTFR\n",
    "    # times=[-1.5, 3.5] #this is for 0.5 sec of padding on each side. Just do 3 seconds as estimate of the full trial. Don't use avg RT because its different across subjects and makes later analysis hard.\n",
    "    times = [-1.5, 0.5] #padd 0.5 secs each side, want 1 seconds before stim onset\n",
    "    print(\"times: \" + str(times))\n",
    "    trials = trial_ieeg(good, \"Stimulus\", times, preload=True)\n",
    "\n",
    "    # print(\"base trials shape is: \" + str(trials._data.shape))\n",
    "    outliers_to_nan(trials, outliers=10)\n",
    "    base = wavelet_scaleogram(trials, n_jobs=1, decim=int(good.info['sfreq'] / 100))\n",
    "    crop_pad(base, \"0.5s\")\n",
    "\n",
    "\n",
    "    print(\"done with base\")\n",
    "\n",
    "    #now do rescale with the concatenated baseline epochs\n",
    "    for event, t in zip((\"Stimulus/i25\", \"Stimulus/c25\"),((-1, 2), (-1, 2))):\n",
    "        times = [None, None]\n",
    "        times[0] = t[0] - 0.5\n",
    "        times[1] = t[1] + 0.5\n",
    "        trials = trial_ieeg(good, event, times, preload=True)\n",
    "        print(f\"trials for {event} is: \" + str(trials._data.shape))\n",
    "        outliers_to_nan(trials, outliers=10)\n",
    "        spec = wavelet_scaleogram(trials, n_jobs=1, decim=int(good.info['sfreq'] / 100))\n",
    "        crop_pad(spec, \"0.5s\")\n",
    "        spec_a = rescale(spec, base, copy=True, mode='ratio').average(\n",
    "                lambda x: np.nanmean(x, axis=0), copy=True)\n",
    "        spec_a._data = np.log10(spec_a._data) * 20\n",
    "        fnames = [os.path.relpath(f, layout.root) for f in good.filenames]\n",
    "        spec_a.info['subject_info']['files'] = tuple(fnames)\n",
    "        spec_a.info['bads'] = good.info['bads']\n",
    "        filename = os.path.join(save_dir, f'{sub}_{event}_1secBeforeStimOnset-tfr.h5')\n",
    "        mne.time_frequency.write_tfrs(filename, spec_a, overwrite=True)\n",
    "        spec_a.save(os.path.join(save_dir, f'{sub}_{event}_1secBeforeStimOnset-avg.fif'), overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### test it as a function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mne.time_frequency\n",
    "import mne\n",
    "from ieeg.io import get_data, raw_from_layout\n",
    "from ieeg.navigate import trial_ieeg, channel_outlier_marker, crop_empty_data, outliers_to_nan\n",
    "from ieeg.calc.scaling import rescale\n",
    "import os\n",
    "from ieeg.timefreq.utils import wavelet_scaleogram, crop_pad\n",
    "import numpy as np\n",
    "from misc_functions import calculate_RTs\n",
    "\n",
    "HOME = os.path.expanduser(\"~\")\n",
    "\n",
    "# get box directory depending on OS\n",
    "if os.name == 'nt': # windows\n",
    "    LAB_root = os.path.join(HOME, \"Box\", \"CoganLab\")\n",
    "else: # mac\n",
    "    LAB_root = os.path.join(HOME, \"Library\", \"CloudStorage\", \"Box-Box\", \"CoganLab\")\n",
    "\n",
    "layout = get_data(\"GlobalLocal\", root=LAB_root)\n",
    "layout = get_data(\"GlobalLocal\", root=LAB_root)\n",
    "\n",
    "def get_wavelet_baseline(inst: mne.io.BaseRaw, base_times: tuple[float, float]):\n",
    "    inst = inst.copy()\n",
    "    inst.load_data()\n",
    "    ch_type = inst.get_channel_types(only_data_chs=True)[0]\n",
    "    inst.set_eeg_reference(ref_channels=\"average\", ch_type=ch_type)\n",
    "\n",
    "    adjusted_base_times = [base_times[0] - 0.5, base_times[1] + 0.5]\n",
    "    trials = trial_ieeg(inst, \"experimentStart\", adjusted_base_times, preload=True)\n",
    "    outliers_to_nan(trials, outliers=10)\n",
    "    base = wavelet_scaleogram(trials, n_jobs=-2, decim=int(inst.info['sfreq'] / 100))\n",
    "    crop_pad(base, \"0.5s\")\n",
    "    del inst\n",
    "    return base\n",
    "\n",
    "def make_wavelets(sub, task, output_name, event=None, times=(-1, 2), base_times=(1,101), LAB_root=None):\n",
    "    \"\"\"\n",
    "    Generate wavelets for a given subject and task using specified events.\n",
    "    \n",
    "    Parameters:\n",
    "    - sub (str): The subject identifier.\n",
    "    - task (str): The task identifier.\n",
    "    - event (str, optional): The event to process. Defaults to None.\n",
    "    - times (tuple, optional): A tuple indicating the start and end times for processing. Defaults to (-1, 2).\n",
    "    - base_times (tuple, optional): A tuple indicating the start and end base times for processing. Defaults to (-0.5, 0).\n",
    "    - LAB_root (str, optional): The root directory for the lab. Will be determined based on OS if not provided. Defaults to None.\n",
    "    - output_name (str): The name for the output files.\n",
    "    \n",
    "    Returns:\n",
    "    - spec_a: Processed data after generating wavelets.\n",
    "    \n",
    "    This function will process the provided events for a given subject and task. \n",
    "    Wavelets will be generated for the events and the results will be saved to output files.\n",
    "    \"\"\"\n",
    "\n",
    "sub='D0057'\n",
    "task=\"GlobalLocal\"\n",
    "output_name=\"experimentStartBase_Stimulus\"\n",
    "event=\"Stimulus\"\n",
    "times=(-1, 2)\n",
    "base_times=(1,101)\n",
    "# Check if LAB_root is provided. If not, determine based on OS.\n",
    "if LAB_root is None:\n",
    "    HOME = os.path.expanduser(\"~\")\n",
    "    if os.name == 'nt': # windows\n",
    "        LAB_root = os.path.join(HOME, \"Box\", \"CoganLab\")\n",
    "    else: # mac\n",
    "        LAB_root = os.path.join(HOME, \"Library\", \"CloudStorage\", \"Box-Box\", \"CoganLab\")\n",
    "\n",
    "layout = get_data(task, root=LAB_root)\n",
    "\n",
    "# Load the data\n",
    "filt = raw_from_layout(layout.derivatives['derivatives/clean'], subject=sub, extension='.edf', desc='clean', preload=False)\n",
    "\n",
    "# Crop raw data to minimize processing time\n",
    "good = crop_empty_data(filt)\n",
    "good.info['bads'] = channel_outlier_marker(good, 3, 2)\n",
    "good.drop_channels(good.info['bads'])\n",
    "good.load_data()\n",
    "\n",
    "ch_type = filt.get_channel_types(only_data_chs=True)[0]\n",
    "good.set_eeg_reference(ref_channels=\"average\", ch_type=ch_type)\n",
    "\n",
    "good.plot()\n",
    "\n",
    "# Epoching and trial outlier removal\n",
    "save_dir = os.path.join(layout.root, 'derivatives', 'spec', 'wavelet', sub)\n",
    "if not os.path.exists(save_dir):\n",
    "    os.makedirs(save_dir)\n",
    "\n",
    "RTs, skipped = calculate_RTs(good)\n",
    "avg_RT = np.median(RTs)\n",
    "\n",
    "# Adjust base_times for padding\n",
    "adjusted_base_times = [base_times[0] - 0.5, base_times[1] + 0.5]\n",
    "\n",
    "base = get_wavelet_baseline(filt, base_times)\n",
    "# # Make stimulus baseline EpochsTFR\n",
    "# trials = trial_ieeg(good, \"Stimulus\", adjusted_base_times, preload=True)\n",
    "# outliers_to_nan(trials, outliers=10)\n",
    "# base = wavelet_scaleogram(trials, n_jobs=-2, decim=int(good.info['sfreq'] / 100))\n",
    "# crop_pad(base, \"0.5s\")\n",
    "\n",
    "all_trials = []\n",
    "\n",
    "# Adjust times for padding\n",
    "times = [times[0] - 0.5, times[1] + 0.5]\n",
    "trials = trial_ieeg(good, event, times, preload=True)\n",
    "all_trials.append(trials)\n",
    "    \n",
    "print(\"all_trials: \", all_trials)\n",
    "\n",
    "# concatenate all trials\n",
    "combined_trials = mne.concatenate_epochs(all_trials)\n",
    "print(\"combined_trials: \", combined_trials)\n",
    "\n",
    "outliers_to_nan(combined_trials, outliers=10)\n",
    "spec = wavelet_scaleogram(combined_trials, n_jobs=-2, decim=int(good.info['sfreq'] / 100))\n",
    "crop_pad(spec, \"0.5s\")\n",
    "spec_a = rescale(spec, base, copy=True, mode='ratio').average(lambda x: np.nanmean(x, axis=0), copy=True)\n",
    "spec_a._data = np.log10(spec_a._data) * 20\n",
    "fnames = [os.path.relpath(f, layout.root) for f in good.filenames]\n",
    "spec_a.info['subject_info']['files'] = tuple(fnames)\n",
    "spec_a.info['bads'] = good.info['bads']\n",
    "filename = os.path.join(save_dir, f'{output_name}-tfr.h5')\n",
    "mne.time_frequency.write_tfrs(filename, spec_a, overwrite=True)\n",
    "spec_a.save(os.path.join(save_dir, f'{output_name}-avg.fif'), overwrite=True)\n",
    "    \n",
    "    # return spec_a  # Optional: Return the processed data for further use.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spec_a = make_wavelets(sub='D0057', task=\"GlobalLocal\", output_name=\"experimentStartBase_Stimulus\", event=\"Stimulus\", times=(-1, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spec_a = make_wavelets(sub='D0063', task=\"GlobalLocal\", output_name=\"repeat\", events_list=[\"/r\"], times=(-1, 2))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ieeg",
   "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.11.7"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}