src/dicarlo_lab_to_nwb/conversion/showcase_notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#  Loading Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### NWB"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Running a conversion\n",
    "First we run one with only the raw data:\n",
    "* Intan raw data.\n",
    "* Behavioral data.\n",
    "* Stimulus data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "data_folder = Path(\"/media/heberto/One Touch/DiCarlo-CN-data-share\")\n",
    "assert data_folder.is_dir(), f\"Data directory not found: {data_folder}\"\n",
    "\n",
    "stimuli_folder = data_folder / \"StimulusSets\" / \"RSVP-domain_transfer\" / \"images\"\n",
    "stub_test = True\n",
    "verbose = True\n",
    "\n",
    "\n",
    "from dicarlo_lab_to_nwb.conversion.data_locator import locate_intan_file_path, locate_mworks_processed_file_path\n",
    "from dicarlo_lab_to_nwb.conversion.convert_session import convert_session_to_nwb\n",
    "\n",
    "session_metadata = {\n",
    "    \"image_set_name\": \"domain-transfer-2023\",\n",
    "    \"session_date\": \"20230215\",\n",
    "    \"session_time\": \"161322\",\n",
    "    \"subject\": \"pico\",\n",
    "}\n",
    "\n",
    "# These two functions is where we encode your data organization structure.\n",
    "intan_file_path = locate_intan_file_path(data_folder=data_folder, **session_metadata)\n",
    "mworks_processed_file_path = locate_mworks_processed_file_path(data_folder=data_folder, **session_metadata)\n",
    "\n",
    "\n",
    "convert_session_to_nwb(\n",
    "    session_metadata=session_metadata,\n",
    "    intan_file_path=intan_file_path,\n",
    "    mworks_processed_file_path=mworks_processed_file_path,\n",
    "    stimuli_folder=stimuli_folder,\n",
    "    stub_test=stub_test,\n",
    "    verbose=verbose,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also include the threshold crossings and the psths data by running the same conversion scripts with enchanced options."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_folder = Path(\"/media/heberto/One Touch/DiCarlo-CN-data-share\")\n",
    "assert data_folder.is_dir(), f\"Data directory not found: {data_folder}\"\n",
    "\n",
    "stimuli_folder = data_folder / \"StimulusSets\" / \"RSVP-domain_transfer\" / \"images\"\n",
    "stub_test = True\n",
    "verbose = True\n",
    "\n",
    "\n",
    "from dicarlo_lab_to_nwb.conversion.data_locator import locate_intan_file_path, locate_mworks_processed_file_path\n",
    "from dicarlo_lab_to_nwb.conversion.convert_session import convert_session_to_nwb\n",
    "\n",
    "session_metadata = {\n",
    "    \"image_set_name\": \"domain-transfer-2023\",\n",
    "    \"session_date\": \"20230215\",\n",
    "    \"session_time\": \"161322\",\n",
    "    \"subject\": \"pico\",\n",
    "}\n",
    "\n",
    "# These two functions is where we encode your data organization structure.\n",
    "intan_file_path = locate_intan_file_path(data_folder=data_folder, **session_metadata)\n",
    "mworks_processed_file_path = locate_mworks_processed_file_path(data_folder=data_folder, **session_metadata)\n",
    "\n",
    "thresholindg_pipeline_kwargs = {\n",
    "    \"f_notch\": 60.0,  # Frequency for the notch filter\n",
    "    \"bandwidth\": 10.0,  # Bandwidth for the notch filter\n",
    "    \"f_low\": 300.0,  # Low cutoff frequency for the bandpass filter\n",
    "    \"f_high\": 6000.0,  # High cutoff frequency for the bandpass filter\n",
    "    \"noise_threshold\": 3,  # Threshold for detection in the thresholding algorithm\n",
    "}\n",
    "\n",
    "# Ten bins starting 200 ms before the stimulus and spanning 400 ms\n",
    "psth_kwargs = {\"bins_span_milliseconds\": 400, \"num_bins\": 10, \"milliseconds_from_event_to_first_bin\": -200.0}\n",
    "\n",
    "convert_session_to_nwb(\n",
    "    session_metadata=session_metadata,\n",
    "    intan_file_path=intan_file_path,\n",
    "    mworks_processed_file_path=mworks_processed_file_path,\n",
    "    stimuli_folder=stimuli_folder,\n",
    "    add_thresholding_events=True,\n",
    "    thresholindg_pipeline_kwargs=thresholindg_pipeline_kwargs,\n",
    "    add_psth=True,\n",
    "    psth_kwargs=psth_kwargs,\n",
    "    stub_test=stub_test,\n",
    "    verbose=verbose,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Loading an NWBFile\n",
    "After running the script on `conversion.convert_session.py` the produced nwbfile can be loaded like this"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from pynwb import NWBHDF5IO\n",
    "\n",
    "# Change this for the location of the NWB files in your system\n",
    "stub_file = False\n",
    "nwb_folder_path = Path.home() / \"conversion_nwb\"   \n",
    "\n",
    "if stub_file:\n",
    "    nwb_folder_path = nwb_folder_path / \"nwb_stub\"\n",
    "    assert nwb_folder_path.is_dir()\n",
    "\n",
    "nwbfile_path = nwb_folder_path / \"pico_20230214_140610.nwb\"\n",
    "assert nwbfile_path.is_file(), f\"{nwbfile_path} does not exist\"\n",
    "\n",
    "\n",
    "io = NWBHDF5IO(nwbfile_path, mode=\"r\")\n",
    "nwbfile = io.read()\n",
    "nwbfile"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Trials Table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nwbfile.trials.to_dataframe().columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "columns = [\n",
    "    \"start_time\",\n",
    "    \"stop_time\",\n",
    "    \"stimulus_presented\",\n",
    "    \"fixation_correct\",\n",
    "    \"stimuli_block_index\",\n",
    "]\n",
    "\n",
    "nwbfile.trials.to_dataframe()[columns].sample(n=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Calculate compression ratio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hdf5_dataset = nwbfile.acquisition[\"ElectricalSeries\"].data\n",
    "size_uncompressed_GiB = hdf5_dataset.nbytes / 1024 ** 3\n",
    "print(f\"Size of the uncompressed ElectricalSeries: {size_uncompressed_GiB:.2f} GiB\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "size_compressed_GiB = hdf5_dataset.id.get_storage_size() / 1024 ** 3\n",
    "print(f\"Size of the compressed ElectricalSeries: {size_compressed_GiB:.2f} GiB\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "compression_ratio = hdf5_dataset.nbytes / hdf5_dataset.id.get_storage_size() \n",
    "print(f\"Compression ratio: {compression_ratio:.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "space_saving = 1 - size_compressed_GiB / size_uncompressed_GiB\n",
    "print(f\"Space saving: {space_saving:.2f} %\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### PSTHs Binned Aligned Spikes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nwbfile.processing[\"ecephys\"][\"BinnedAlignedSpikesStimulusID0\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nwbfile.processing[\"ecephys\"][\"BinnedAlignedSpikesStimulusID0\"].data[:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nwbfile.processing[\"ecephys\"][\"BinnedAlignedSpikesStimulusID0\"].event_timestamps[:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Display stimuli"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "an_image = nwbfile.stimulus[\"stimuli\"].images[\"im0\"]\n",
    "\n",
    "an_image_data = an_image.data[:]\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.imshow(an_image_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Electrode Table "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "columns = [\n",
    "    \"channel_name\",\n",
    "    \"probe\",\n",
    "    \"rel_x\",\n",
    "    \"rel_y\",\n",
    "    \"electrode_impedance_magnitude\",\n",
    "    \"electrode_impedance_phase\",\n",
    "]\n",
    "nwbfile.electrodes.to_dataframe()[columns].sample(n=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Units Table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nwbfile.units.to_dataframe().sample(n=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Amplifiers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is an example of how to load data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from dicarlo_lab_to_nwb.conversion.data_locator import locate_intan_file_path\n",
    "\n",
    "data_folder = Path(\"/media/heberto/One Touch/DiCarlo-CN-data-share\")\n",
    "image_set_name = \"domain-transfer-2023\"\n",
    "subject = \"pico\"\n",
    "session_date = \"20230214\"\n",
    "session_time = \"140610\"\n",
    "\n",
    "intan_file_path = locate_intan_file_path(\n",
    "    data_folder=data_folder,\n",
    "    image_set_name=image_set_name,\n",
    "    subject=subject,\n",
    "    session_date=session_date,\n",
    "    session_time=session_time,\n",
    ")\n",
    "\n",
    "from spikeinterface.extractors import IntanRecordingExtractor\n",
    "recording = IntanRecordingExtractor(\n",
    "    file_path=intan_file_path,  \n",
    "    stream_name=\"RHD2000 amplifier channel\",\n",
    "    all_annotations=True,\n",
    "    ignore_integrity_checks=False,\n",
    ")\n",
    "recording"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This particular example has timestamps discontinuities, to load the data regardless we set the parameter `ignore_integrity_checks=True`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from spikeinterface.extractors import IntanRecordingExtractor\n",
    "\n",
    "recording = IntanRecordingExtractor(\n",
    "    file_path=intan_file_path,\n",
    "    stream_name=\"RHD2000 amplifier channel\",\n",
    "    all_annotations=True,  # the .rhd file\n",
    "    ignore_integrity_checks=True,\n",
    ")\n",
    "recording"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Auxiliary input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "recording_auxiliary_input = IntanRecordingExtractor(\n",
    "    file_path=intan_file_path,\n",
    "    stream_name=\"RHD2000 auxiliary input channel\",\n",
    "    all_annotations=True,\n",
    "    ignore_integrity_checks=True,\n",
    ")\n",
    "\n",
    "recording_auxiliary_input"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ADC input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "recording_adc_input = IntanRecordingExtractor(\n",
    "    file_path=intan_file_path,\n",
    "    stream_name=\"USB board ADC input channel\",\n",
    "    all_annotations=True,\n",
    "    ignore_integrity_checks=True,\n",
    ")\n",
    "\n",
    "recording_adc_input"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Digital channel \n",
    "Requires neo version from github https://github.com/NeuralEnsemble/python-neo/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "recording_digital = IntanRecordingExtractor(\n",
    "    file_path=intan_file_path,\n",
    "    stream_name=\"USB board digital input channel\",\n",
    "    all_annotations=True,\n",
    "    ignore_integrity_checks=True,\n",
    ")\n",
    "\n",
    "recording_digital"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading the probe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dicarlo_lab_to_nwb.conversion.probe import build_probe_group\n",
    "from dicarlo_lab_to_nwb.conversion.data_locator import locate_intan_file_path\n",
    "from spikeinterface.extractors import IntanRecordingExtractor\n",
    "\n",
    "data_folder = Path(\"/media/heberto/One Touch/DiCarlo-CN-data-share\")\n",
    "image_set_name = \"domain-transfer-2023\"\n",
    "subject = \"pico\"\n",
    "session_date = \"20230214\"\n",
    "session_time = \"140610\"\n",
    "\n",
    "\n",
    "intan_file_path = locate_intan_file_path(\n",
    "    data_folder=data_folder,\n",
    "    image_set_name=image_set_name,\n",
    "    subject=subject,\n",
    "    session_date=session_date,\n",
    "    session_time=session_time,\n",
    ")\n",
    "\n",
    "\n",
    "stream_name = \"RHD2000 amplifier channel\"\n",
    "recording = IntanRecordingExtractor(\n",
    "    file_path=intan_file_path,\n",
    "    stream_name=stream_name,\n",
    "    ignore_integrity_checks=True,\n",
    "    all_annotations=True,\n",
    ")\n",
    "\n",
    "\n",
    "probe_group = build_probe_group(recording=recording)\n",
    "\n",
    "\n",
    "from probeinterface.plotting import plot_probe\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "\n",
    "fig = plt.figure(figsize=(16, 8))\n",
    "ax = fig.add_subplot(111)\n",
    "\n",
    "probe = probe_group.probes[0]\n",
    "channel_ids = recording.get_channel_ids()\n",
    "corresponding_channel_ids = [channel_ids[i] for i in probe.device_channel_indices]\n",
    "\n",
    "text_on_contact = np.asarray(corresponding_channel_ids)\n",
    "\n",
    "plot_probe(probe=probe, ax=ax, with_contact_id=True, text_on_contact=text_on_contact)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from probeinterface.plotting import plot_probe_group\n",
    "\n",
    "fig = plt.figure(figsize=(16, 8))\n",
    "ax = fig.add_subplot(111)\n",
    "\n",
    "\n",
    "plot_probe_group(probe_group, ax=ax, same_axes=True, with_contact_id=False)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sorting Pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run a sorting pipeline we need a recording with a geometry attached."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from spikeinterface.extractors import IntanRecordingExtractor\n",
    "from spikeinterface.sorters import run_sorter_by_property\n",
    "\n",
    "\n",
    "from dicarlo_lab_to_nwb.conversion.data_locator import locate_intan_file_path\n",
    "from dicarlo_lab_to_nwb.conversion.probe import attach_probe_to_recording\n",
    "\n",
    "data_folder = Path(\"/media/heberto/One Touch/DiCarlo-CN-data-share\")\n",
    "image_set_name = \"domain-transfer-2023\"\n",
    "subject = \"pico\"\n",
    "session_date = \"20230214\"\n",
    "session_time = \"140610\"\n",
    "\n",
    "\n",
    "intan_file_path = locate_intan_file_path(\n",
    "    data_folder=data_folder,\n",
    "    image_set_name=image_set_name,\n",
    "    subject=subject,\n",
    "    session_date=session_date,\n",
    "    session_time=session_time,\n",
    ")\n",
    "\n",
    "\n",
    "stream_name = \"RHD2000 amplifier channel\"\n",
    "recording = IntanRecordingExtractor(\n",
    "    file_path=intan_file_path,\n",
    "    stream_name=stream_name,\n",
    "    ignore_integrity_checks=True,\n",
    "    all_annotations=True,\n",
    ")\n",
    "\n",
    "\n",
    "attach_probe_to_recording(recording=recording)\n",
    "recording"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Most sorters have been designed with high density probes in mind. They will work with a single channel probe, but the results may not be as good as some units might be supressed by the spatial regularization.\n",
    "\n",
    "Because of this we performed sorting in two ways so you can compare the results:\n",
    "\n",
    "1. We do one sorting per probe\n",
    "2. We do one sorting per channel to avoid interference of the spatial regularization\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Performing a sorting per probe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from spikeinterface.core import load_extractor\n",
    "\n",
    "sorting_folder = Path(\"./sorting_done\")\n",
    "overwrite = False\n",
    "\n",
    "if sorting_folder.exists() and not overwrite:\n",
    "    sorting = load_extractor(sorting_folder)\n",
    "else:\n",
    "    sorting = run_sorter_by_property(\n",
    "        sorter_name=\"kilosort2\",\n",
    "        recording=recording,\n",
    "        folder=\"./sorting_folder_probe\",\n",
    "        grouping_property=\"probe\",\n",
    "        docker_image=True,\n",
    "    )\n",
    "\n",
    "    sorting.save(folder=sorting_folder)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sorting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from spikeinterface.core import create_sorting_analyzer\n",
    "\n",
    "\n",
    "sorting_analyzer = create_sorting_analyzer(sorting=sorting, recording=recording)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Performing a sorting per channel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sorting = run_sorter_by_property(\n",
    "    sorter_name=\"kilosort3\",\n",
    "    recording=recording,\n",
    "    folder=\"./sorting_folder_per_channel\",\n",
    "    grouping_property=\"channel_names\",\n",
    "    docker_image=True,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from spikeinterface.sorters import available_sorters\n",
    "\n",
    "available_sorters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from spikeinterface.core import load_extractor\n",
    "\n",
    "sorting_folder = Path(\"./sorting_done_per_channel\")\n",
    "overwrite = False\n",
    "\n",
    "if sorting_folder.exists() and not overwrite:\n",
    "    sorting = load_extractor(sorting_folder)\n",
    "else:\n",
    "    sorting = run_sorter_by_property(\n",
    "        sorter_name=\"tridesclous\",\n",
    "        recording=recording,\n",
    "        folder=\"./sorting_folder_per_channel\",\n",
    "        grouping_property=\"channel_names\",\n",
    "        docker_image=True,\n",
    "    )\n",
    "\n",
    "    sorting.save(folder=sorting_folder)\n",
    "    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Peak Detection Pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Artificial data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import spikeinterface.widgets as sw\n",
    "\n",
    "from spikeinterface.core.generate import generate_ground_truth_recording\n",
    "\n",
    "\n",
    "recording, sorting = generate_ground_truth_recording(num_channels=4, num_units=1, durations=[1], seed=0)\n",
    "\n",
    "\n",
    "w_ts = sw.plot_traces(recording, time_range=(0, 1))\n",
    "w_rs = sw.plot_rasters(sorting, time_range=(0, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from dicarlo_lab_to_nwb.conversion.pipeline import thresholding_pipeline\n",
    "\n",
    "\n",
    "job_kwargs = dict(n_jobs=1, progress_bar=True, chunk_duration=1.0)\n",
    "noise_threshold = 3  # The number of standard deviations for peak detection\n",
    "\n",
    "spike_times_per_channel = thresholding_pipeline(\n",
    "    recording=recording,\n",
    "    noise_threshold=noise_threshold,\n",
    "    job_kwargs=job_kwargs,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sorting.get_unit_spike_train(0, return_times=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spike_times_per_channel[0][0] * 1000.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Intan Recording data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import spikeinterface.widgets as sw\n",
    "from dicarlo_lab_to_nwb.conversion.data_locator import locate_intan_file_path\n",
    "\n",
    "from spikeinterface.extractors import IntanRecordingExtractor\n",
    "\n",
    "\n",
    "data_folder = Path(\"/media/heberto/One Touch/DiCarlo-CN-data-share\")\n",
    "image_set_name = \"domain-transfer-2023\"\n",
    "subject = \"pico\"\n",
    "session_date = \"20230214\"\n",
    "session_time = \"140610\"\n",
    "\n",
    "\n",
    "intan_file_path = locate_intan_file_path(\n",
    "    data_folder=data_folder,\n",
    "    image_set_name=image_set_name,\n",
    "    subject=subject,\n",
    "    session_date=session_date,\n",
    "    session_time=session_time,\n",
    ")\n",
    "\n",
    "\n",
    "\n",
    "recording = IntanRecordingExtractor(\n",
    "    file_path=intan_file_path,\n",
    "    stream_name=\"RHD2000 amplifier channel\",\n",
    "    all_annotations=True,\n",
    "    ignore_integrity_checks=True,\n",
    ")\n",
    "\n",
    "# If you want to select only one channel\n",
    "channel_ids = recording.get_channel_ids()[0:1]\n",
    "single_channel_recording = recording.select_channels(channel_ids=channel_ids)\n",
    "w_ts = sw.plot_traces(single_channel_recording, time_range=(0, 1), return_scaled=True)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Preprocess"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dicarlo_lab_to_nwb.conversion.pipeline import DiCarloBandPass, DiCarloNotch\n",
    "\n",
    "\n",
    "f_notch = 60  # Hz\n",
    "bandwidth = 10\n",
    "f_low = 300.0\n",
    "f_high = 6000.0\n",
    "\n",
    "vectorized = True \n",
    "notched_recording = DiCarloNotch(single_channel_recording, f_notch=f_notch, bandwidth=bandwidth, vectorized=vectorized)\n",
    "preprocessed_recording = DiCarloBandPass(notched_recording, f_low=f_low, f_high=f_high, vectorized=vectorized)\n",
    "\n",
    "# For this instance each array 96 channels, 400 micrometes apart\n",
    "w_ts = sw.plot_traces(preprocessed_recording, time_range=(0, 1), return_scaled=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Run the peak detection on a short portion of the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dicarlo_lab_to_nwb.conversion.pipeline import thresholding_pipeline\n",
    "\n",
    "noise_threshold = 3  # The number of standard deviations for peak detection\n",
    "\n",
    "start_time = 0\n",
    "end_time = 10.0\n",
    "\n",
    "preprocessed_recording = preprocessed_recording.time_slice(start_time=start_time, end_time=end_time)\n",
    "\n",
    "spike_times_per_channel = thresholding_pipeline(\n",
    "    recording=preprocessed_recording,\n",
    "    noise_threshold=noise_threshold,\n",
    ")\n",
    "\n",
    "spike_times_per_channel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Everything can be wrapped up in a couple of lines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from spikeinterface.extractors import IntanRecordingExtractor\n",
    "from dicarlo_lab_to_nwb.conversion.data_locator import locate_intan_file_path\n",
    "\n",
    "data_folder = Path(\"/media/heberto/One Touch/DiCarlo-CN-data-share\")\n",
    "image_set_name = \"domain-transfer-2023\"\n",
    "subject = \"pico\"\n",
    "session_date = \"20230214\"\n",
    "session_time = \"140610\"\n",
    "\n",
    "\n",
    "intan_file_path = locate_intan_file_path(\n",
    "    data_folder=data_folder,\n",
    "    image_set_name=image_set_name,\n",
    "    subject=subject,\n",
    "    session_date=session_date,\n",
    "    session_time=session_time,\n",
    ")\n",
    "\n",
    "stream_name = \"RHD2000 amplifier channel\"\n",
    "recording = IntanRecordingExtractor(\n",
    "    file_path=intan_file_path,\n",
    "    stream_name=stream_name,\n",
    "    ignore_integrity_checks=True,\n",
    "    all_annotations=True,\n",
    ")\n",
    "\n",
    "from dicarlo_lab_to_nwb.conversion.pipeline import thresholding_pipeline\n",
    "# Parameters of the pipeline\n",
    "f_notch = 60  # Hz\n",
    "bandwidth = 10\n",
    "f_low = 300.0\n",
    "f_high = 6000.0\n",
    "noise_threshold = 3  # The number of standard deviations for peak detection\n",
    "\n",
    "stub_test = False   \n",
    "if stub_test:\n",
    "    recording = recording.time_slice(start_time=0, end_time=60.0 * 5)\n",
    "\n",
    "spike_times_per_channel = thresholding_pipeline(\n",
    "    recording=recording,\n",
    "    f_notch=f_notch,\n",
    "    bandwidth=bandwidth,\n",
    "    f_low=f_low,\n",
    "    f_high=f_high,\n",
    "    noise_threshold=noise_threshold,\n",
    "    verbose=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output of the pipeline is a dictionary whose keys are the channel ids and the values are the times (in seconds) at which the threshold was crossed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spike_times_per_channel[\"A-000\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### The algorithm can be run from NWB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from dicarlo_lab_to_nwb.conversion.pipeline import thresholding_pipeline\n",
    "\n",
    "\n",
    "# Change this for the location of the NWB files in your system\n",
    "nwb_folder_path = Path.home() / \"conversion_nwb\"   \n",
    "\n",
    "\n",
    "nwbfile_path = nwb_folder_path / \"pico_20230214_140610.nwb\"\n",
    "\n",
    "from spikeinterface.extractors import NwbRecordingExtractor\n",
    "\n",
    "recording = NwbRecordingExtractor(file_path=nwbfile_path)\n",
    "\n",
    "# Parameters of the pipeline\n",
    "f_notch = 60  # Hz\n",
    "bandwidth = 10\n",
    "f_low = 300.0\n",
    "f_high = 6000.0\n",
    "noise_threshold = 3  # The number of standard deviations for peak detection\n",
    "\n",
    "stub_test = False   \n",
    "if stub_test:\n",
    "    recording = recording.time_slice(start_time=0, end_time=60.0 * 5)\n",
    "\n",
    "spike_times_per_channel = thresholding_pipeline(\n",
    "    recording=recording,\n",
    "    f_notch=f_notch,\n",
    "    bandwidth=bandwidth,\n",
    "    f_low=f_low,\n",
    "    f_high=f_high,\n",
    "    noise_threshold=noise_threshold,\n",
    "    verbose=True,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Calculating PSTH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload \n",
    "from pathlib import Path\n",
    "\n",
    "from spikeinterface.extractors import IntanRecordingExtractor\n",
    "from dicarlo_lab_to_nwb.conversion.pipeline import thresholding_pipeline\n",
    "from dicarlo_lab_to_nwb.conversion.data_locator import locate_intan_file_path\n",
    "from dicarlo_lab_to_nwb.conversion.probe import attach_probe_to_recording\n",
    "\n",
    "data_folder = Path(\"/media/heberto/One Touch/DiCarlo-CN-data-share\")\n",
    "image_set_name = \"domain-transfer-2023\"\n",
    "subject = \"pico\"\n",
    "session_date = \"20230214\"\n",
    "session_time = \"140610\"\n",
    "\n",
    "# Parameters of the pipeline\n",
    "f_notch = 60  # Hz\n",
    "bandwidth = 10\n",
    "f_low = 300.0\n",
    "f_high = 6000.0\n",
    "noise_threshold = 3  # The number of standard deviations for peak detection\n",
    "\n",
    "intan_file_path = locate_intan_file_path(\n",
    "    data_folder=data_folder,\n",
    "    image_set_name=image_set_name,\n",
    "    subject=subject,\n",
    "    session_date=session_date,\n",
    "    session_time=session_time,\n",
    ")\n",
    "\n",
    "\n",
    "stream_name = \"RHD2000 amplifier channel\"\n",
    "recording = IntanRecordingExtractor(\n",
    "    file_path=intan_file_path,\n",
    "    stream_name=stream_name,\n",
    "    ignore_integrity_checks=True,\n",
    "    all_annotations=True,\n",
    ")\n",
    "\n",
    "\n",
    "attach_probe_to_recording(recording=recording)\n",
    "chunk_duration = 10.0  # 10 seconds\n",
    "job_kwargs = dict(n_jobs=-1, progress_bar=True, chunk_duration=chunk_duration)\n",
    "verbose = True \n",
    "\n",
    "dict_of_recordings = recording.split_by(property=\"probe\", outputs=\"dict\")\n",
    "dict_of_spikes_times_per_channel = {}\n",
    "\n",
    "for probe_name, recording in dict_of_recordings.items():\n",
    "    spikes_times_per_channel = thresholding_pipeline(\n",
    "        recording=recording,\n",
    "        f_notch=f_notch,\n",
    "        bandwidth=bandwidth,\n",
    "        f_low=f_low,\n",
    "        f_high=f_high,\n",
    "        noise_threshold=noise_threshold,\n",
    "        job_kwargs=job_kwargs,\n",
    "        verbose=verbose\n",
    "    )\n",
    "    \n",
    "    dict_of_spikes_times_per_channel[probe_name] = spikes_times_per_channel\n",
    "\n",
    "# We merge all the dictionaries\n",
    "dict_of_spikes_times = {key: value for d in dict_of_spikes_times_per_channel.values() for key, value in d.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from dicarlo_lab_to_nwb.conversion.data_locator import locate_mworks_processed_file_path\n",
    "\n",
    "mworks_processed_file_path = locate_mworks_processed_file_path(\n",
    "    data_folder=data_folder,\n",
    "    image_set_name=image_set_name,\n",
    "    subject=subject,\n",
    "    session_date=session_date,\n",
    "    session_time=session_time,\n",
    ")\n",
    "\n",
    "\n",
    "mworks_processed_file_path = Path(mworks_processed_file_path)\n",
    "dtype = {\"stimulus_presented\": np.uint32, \"fixation_correct\": bool}\n",
    "mwkorks_df = pd.read_csv(mworks_processed_file_path, dtype=dtype)\n",
    "ground_truth_time_column = \"samp_on_us\"\n",
    "stimuli_presentation_times_seconds = mwkorks_df[ground_truth_time_column] / 1e6\n",
    "stimuli_presentation_id = mwkorks_df[\"stimulus_presented\"]\n",
    "stimuli_ids = stimuli_presentation_id.unique()\n",
    "stimuli_ids_sorted = sorted(stimuli_ids) \n",
    "# Sort the stimuli by file-name\n",
    "\n",
    "stimuli_presentation_times_dict = {\n",
    "    stimulus_id: stimuli_presentation_times_seconds[stimuli_presentation_id == stimulus_id].values for stimulus_id in stimuli_ids_sorted\n",
    "}\n",
    "\n",
    "# Be sure that the list is sorted by unit/site name.\n",
    "unit_ids = list(dict_of_spikes_times.keys())\n",
    "unit_ids_sorted = sorted(unit_ids)\n",
    "spike_times_list = [dict_of_spikes_times[id] for id in unit_ids_sorted]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dicarlo_lab_to_nwb.conversion.psth import calculate_event_psth\n",
    "\n",
    "number_of_bins = 10\n",
    "bins_span_milliseconds = 400.0\n",
    "bin_width_in_milliseconds = bins_span_milliseconds / number_of_bins\n",
    "#This means the first bin starts 200 ms before the image presentation\n",
    "milliseconds_from_event_to_first_bin = -200.0  # \n",
    "max_repetitions = stimuli_presentation_id.value_counts().max()\n",
    "\n",
    "# Let's calculate the PSTH for a single stimuli\n",
    "a_stimuli = stimuli_ids_sorted[0]\n",
    "stimulus_presentation_times = stimuli_presentation_times_dict[a_stimuli]\n",
    "psth_per_stimuli = calculate_event_psth(\n",
    "    spike_times_list=spike_times_list,\n",
    "    event_times_seconds=stimulus_presentation_times,\n",
    "    bin_width_in_milliseconds=bin_width_in_milliseconds,\n",
    "    number_of_bins = number_of_bins,\n",
    "    milliseconds_from_event_to_first_bin=milliseconds_from_event_to_first_bin,\n",
    "    number_of_events=max_repetitions,\n",
    ")\n",
    "\n",
    "psth_per_stimuli[1, ...]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dicarlo_lab_to_nwb.conversion.psth import calculate_event_psth_numpy_naive\n",
    "\n",
    "stimulus_presentation_times = stimuli_presentation_times_dict[a_stimuli]\n",
    "psth_per_stimuli = calculate_event_psth_numpy_naive(\n",
    "    spike_times_list=spike_times_list,\n",
    "    event_times_seconds=stimulus_presentation_times,\n",
    "    bin_width_in_milliseconds=bin_width_in_milliseconds,\n",
    "    number_of_bins = number_of_bins,\n",
    "    milliseconds_from_event_to_first_bin=milliseconds_from_event_to_first_bin,\n",
    "    number_of_events=max_repetitions,\n",
    ")\n",
    "\n",
    "psth_per_stimuli[1, ...]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Agregate psth for all stimuli in session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.auto import tqdm\n",
    "from dicarlo_lab_to_nwb.conversion.psth import calculate_event_psth\n",
    "import time\n",
    "\n",
    "number_of_units = len(spike_times_list)\n",
    "number_of_stimuli = len(stimuli_presentation_times_dict)\n",
    "\n",
    "session_psth = np.full(\n",
    "    shape=(number_of_units, number_of_stimuli, max_repetitions, number_of_bins), fill_value=np.nan\n",
    ")\n",
    "desc = \"Calculating PSTH for stimuli\"\n",
    "\n",
    "time_start = time.time()\n",
    "for stimulus_index, stimuli_id in enumerate(tqdm(stimuli_ids_sorted, desc=desc, unit=\" stimuli processed\")):\n",
    "    stimulus_presentation_times = stimuli_presentation_times_dict[stimuli_id]\n",
    "    psth_per_stimuli = calculate_event_psth(\n",
    "        spike_times_list=spike_times_list,\n",
    "        event_times_seconds=stimulus_presentation_times, \n",
    "        bin_width_in_milliseconds=bin_width_in_milliseconds,\n",
    "        number_of_bins=number_of_bins,\n",
    "        milliseconds_from_event_to_first_bin=milliseconds_from_event_to_first_bin,\n",
    "        number_of_events=max_repetitions,\n",
    "    )\n",
    "    session_psth[:, stimulus_index, :, :] = psth_per_stimuli\n",
    "    \n",
    "session_psth_numba = session_psth\n",
    "\n",
    "time_stop = time.time()\n",
    "time_numba = time_stop - time_start\n",
    "print(f\"Time elapsed: {time_numba:2.2f} seconds\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"session_psth.shape: {session_psth.shape}, {number_of_stimuli=}, {max_repetitions=}, {number_of_bins=}, {number_of_units=}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To transform the data into the DiCarlo lab format a simple transport suffices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session_psth.transpose((1, 2, 3, 0)).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To compare to a baseline this is the naive numpy implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "from dicarlo_lab_to_nwb.conversion.psth import calculate_event_psth_numpy_naive\n",
    "import time \n",
    "\n",
    "number_of_units = len(spike_times_list)\n",
    "number_of_stimuli = len(stimuli_presentation_times_dict)\n",
    "\n",
    "session_psth = np.full(\n",
    "    shape=(number_of_units, number_of_stimuli, max_repetitions, number_of_bins), fill_value=np.nan\n",
    ")\n",
    "\n",
    "start = time.time()\n",
    "for stimulus_index, (stimulus_id, stimulus_times) in enumerate(tqdm(stimuli_presentation_times_dict.items(), desc=\"Processing Stimuli\")):\n",
    "    psth_per_stimuli = calculate_event_psth_numpy_naive(\n",
    "        spike_times_list=spike_times_list,\n",
    "        event_times_seconds=stimulus_times,  # make sure this is correct\n",
    "        bin_width_in_milliseconds=bin_width_in_milliseconds,\n",
    "        number_of_bins=number_of_bins,\n",
    "        milliseconds_from_event_to_first_bin=milliseconds_from_event_to_first_bin,\n",
    "        number_of_events=max_repetitions,\n",
    "    )\n",
    "    session_psth[:, stimulus_index, :, :] = psth_per_stimuli\n",
    "    \n",
    "session_psth_naive = session_psth\n",
    "\n",
    "time_stop = time.time()\n",
    "time_naive = time_stop - time_start\n",
    "print(f\"Time elapsed: {time_naive:2.2f} seconds\")\n",
    "time_naive_minutes = time_naive / 60.0\n",
    "print(f\"Time elapsed: {time_naive_minutes:2.2f} minutes\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.allclose(session_psth_numba, session_psth_naive, equal_nan=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "On my local machine this is ~10x faster than the naive numpy implementation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculating PSTH from NWBFile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from pynwb import NWBHDF5IO\n",
    "\n",
    "nwb_folder_path = Path.home() / \"conversion_nwb\"   \n",
    "\n",
    "nwbfile_path = nwb_folder_path / \"pico_20230214_140610.nwb\"\n",
    "assert nwbfile_path.is_file(), f\"{nwbfile_path} does not exist\"\n",
    "\n",
    "\n",
    "io = NWBHDF5IO(nwbfile_path, mode=\"r\")\n",
    "nwbfile = io.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dicarlo_lab_to_nwb.conversion.psth import build_psth_from_nwbfile\n",
    "\n",
    "number_of_bins = 10\n",
    "bins_span_milliseconds = 400.0\n",
    "bin_width_in_milliseconds = bins_span_milliseconds / number_of_bins\n",
    "#This means the first bin starts 200 ms before the image presentation\n",
    "milliseconds_from_event_to_first_bin = -200.0  \n",
    "\n",
    "psth_dict, stimuli_presentation_times_dict = build_psth_from_nwbfile(\n",
    "    nwbfile=nwbfile,\n",
    "    bin_width_in_milliseconds=bin_width_in_milliseconds,\n",
    "    number_of_bins=number_of_bins,\n",
    "    milliseconds_from_event_to_first_bin=milliseconds_from_event_to_first_bin,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Aggregating PSTHs from multiple sesssions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from pynwb import NWBHDF5IO\n",
    "import numpy as np\n",
    "\n",
    "# Change this for the location of the NWB files in your system\n",
    "nwb_folder_path = Path.home() / \"conversion_nwb\"   \n",
    "\n",
    "# Here we could filter for a type of experiment using the naming convention, the foler structure or the metadata in the file\n",
    "file_names_available = [path.name for path in nwb_folder_path.iterdir() if path.is_file() and path.suffix == \".nwb\"]\n",
    "\n",
    "nwbfile_list = []\n",
    "for file_name in file_names_available:\n",
    "    nwbfile_file_path = nwb_folder_path / file_name\n",
    "    io = NWBHDF5IO(nwbfile_file_path, 'r')\n",
    "    nwbfile = io.read()\n",
    "    nwbfile_list.append(nwbfile)\n",
    "    \n",
    "    \n",
    "\n",
    "# We have a list per experimental session\n",
    "psth_dict_list = []\n",
    "for nwbfile in nwbfile_list:\n",
    "    interfaces = nwbfile.processing[\"ecephys\"].data_interfaces.values()\n",
    "    is_binned_spikes = lambda interface: interface.data_type == \"BinnedAlignedSpikes\"\n",
    "    valid_interfaces = [interface for interface in interfaces if is_binned_spikes(interface)]\n",
    "    psth_dict = {interface.name: interface.data for interface in valid_interfaces}\n",
    "    psth_dict_list.append(psth_dict)\n",
    "    \n",
    "\n",
    "all_stimuli = set().union(*[psth_dict.keys() for psth_dict in psth_dict_list])\n",
    "# Aggregate psth per stimuli over NWBFiles\n",
    "psth_per_stimuli_dict = {}\n",
    "for stimuli_id in all_stimuli:\n",
    "    stimuli_psth_list = (psth_dict.get(stimuli_id, None) for psth_dict in psth_dict_list)\n",
    "    stimuli_psth_list = [psth for psth in stimuli_psth_list if psth is not None]\n",
    "    stimuli_psth_aggregated = np.concatenate(stimuli_psth_list, axis=1)\n",
    "    psth_per_stimuli_dict[stimuli_id] = stimuli_psth_aggregated\n",
    "\n",
    "# Calculate the max number of repetitions for any stimuli    \n",
    "max_repetitions = max([psth.shape[1] for psth in psth_per_stimuli_dict.values()])\n",
    "num_units = list(psth_per_stimuli_dict.values())[0].shape[0]\n",
    "num_bins = list(psth_per_stimuli_dict.values())[0].shape[2]\n",
    "di_carlo_shape = (len(all_stimuli), max_repetitions, num_bins, num_units)\n",
    "\n",
    "# Coerce to di carlo format shape and fill with nan\n",
    "aggregated_psth = np.full(shape=di_carlo_shape, fill_value=np.nan)\n",
    "for stimuli_index, stimuli_psth in enumerate(psth_per_stimuli_dict.values()):\n",
    "    psth_di_carlo = stimuli_psth.transpose(1, 2, 0)\n",
    "    events_per_stimuli = stimuli_psth.shape[1]\n",
    "    aggregated_psth[stimuli_index, :events_per_stimuli, ...] = psth_di_carlo\n",
    "    \n",
    "aggregated_psth.shape\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "work",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}