# install cell skipped during CI (deps preinstalled into system Python)

import h5py
import remfile
from dandi.dandiapi import DandiAPIClient
from pynwb import NWBHDF5IO

# Connect to DANDI and get the dandiset
dandiset_id = "000409"
client = DandiAPIClient()
dandiset = client.get_dandiset(dandiset_id, "draft")

# =============================================================================
# Session EIDs for NEW format files (desc-raw / desc-processed)
# =============================================================================

# Complete session with 2 probes - Session 1 (NYU-39, 2021-05-10, angelakilab)
# - Full data: all videos, pose estimation, spike sorting for both probes
TWO_PROBE_SESSION_EID_1 = "6ed57216-498d-48a6-b48b-a243a34710ea"

# Complete session with 2 probes - Session 2 (NYU-39, 2021-05-11, angelakilab)
# - Full data: all videos, pose estimation, spike sorting for both probes
TWO_PROBE_SESSION_EID_2 = "35ed605c-1a1a-47b1-86ff-2b56144f55af"

# Complete session with 1 probe (NYU-46, 2021-06-25, angelakilab)  
# - Full data: all videos, pose estimation, spike sorting for probe01
ONE_PROBE_SESSION_EID = "64e3fb86-928c-4079-865c-b364205b502e"

# Choose which session to use
session_eid = ONE_PROBE_SESSION_EID  # Change to TWO_PROBE_SESSION_EID_2 or ONE_PROBE_SESSION_EID
session_eid = "78b4fff5-c5ec-44d9-b5f9-d59493063f00"
# session_eid = "6713a4a7-faed-4df2-acab-ee4e63326f8"
# session_eid = "196a2adf-ff83-49b2-823a-33f990049c2"

# =============================================================================
# Fetch assets by EID
# =============================================================================

# First, filter asse    ts by EID
session_assets = [asset for asset in dandiset.get_assets() if session_eid in asset.path]

# Then, extract raw and processed files
raw_asset = next((asset for asset in session_assets if "desc-raw" in asset.path), None)
processed_asset = next((asset for asset in session_assets if "desc-processed" in asset.path), None)

print(f"Session EID: {session_eid}")
print(f"\nRaw file:       {raw_asset.path if raw_asset else 'Not found'}")
print(f"Processed file: {processed_asset.path if processed_asset else 'Not found'}")

Session EID: 78b4fff5-c5ec-44d9-b5f9-d59493063f00

Raw file:       sub-KS091/sub-KS091_ses-78b4fff5-c5ec-44d9-b5f9-d59493063f00_desc-raw_ecephys.nwb
Processed file: sub-KS091/sub-KS091_ses-78b4fff5-c5ec-44d9-b5f9-d59493063f00_desc-processed_behavior+ecephys.nwb

s3_url = raw_asset.get_content_url(follow_redirects=1, strip_query=False)
file_system = remfile.File(s3_url)
file = h5py.File(file_system, mode="r")

io = NWBHDF5IO(file=file)
nwbfile_raw = io.read()

nwbfile_raw.session_start_time

datetime.datetime(2022, 7, 4, 14, 30, 15, 49000, tzinfo=tzoffset(None, 3600))

devices = nwbfile_raw.devices
probe_devices = [devices[name] for name in devices if "Probe" in name]
for probe_device in probe_devices:
    print(f"{probe_device.name}: {probe_device.serial_number}")

Probe00: 18005102151
Probe01: 18194814141

nwbfile_raw

ap_series = {name: nwbfile_raw.acquisition[name] for name in nwbfile_raw.acquisition if name.endswith("AP")}
for name, series in ap_series.items():
    print(f"{name}: {series.data.shape}")

# Display the first AP series
first_ap = next(iter(ap_series.values()))
first_ap

ElectricalSeriesProbe00AP: (164362247, 384)
ElectricalSeriesProbe01AP: (164362400, 384)

lf_series = {name: nwbfile_raw.acquisition[name] for name in nwbfile_raw.acquisition if name.endswith("LF")}
for name, series in lf_series.items():
    print(f"{name}: {series.data.shape}")

# Display the first LF series
first_lf = next(iter(lf_series.values()))
first_lf

ElectricalSeriesProbe00LF: (13696854, 384)
ElectricalSeriesProbe01LF: (13696867, 384)

first_ap.electrodes.to_dataframe()

electrodes = nwbfile_raw.electrodes
electrodes

electrodes_df = electrodes.to_dataframe()
electrodes_df.columns

Index(['location', 'group', 'electrode_name', 'contact_shapes', 'channel_name',
       'rel_x', 'rel_y', 'probe_name', 'contact_ids', 'shank_ids', 'adc_group',
       'adc_sample_order', 'group_name', 'x', 'y', 'z', 'imp', 'filtering'],
      dtype='object')

columns = ["probe_name", "location", "x", "y", "z", "rel_x", "rel_y"]
electrodes_df[columns].sample(n=5)

localization = nwbfile_raw.lab_meta_data["localization"]
localization

ccf_table_df = localization.anatomical_coordinates_tables["ElectrodesCCFv3"].to_dataframe()
ibl_bregma_table_df = localization.anatomical_coordinates_tables["ElectrodesIBLBregma"].to_dataframe()

ccf_table_df.columns

Index(['x', 'y', 'z', 'localized_entity', 'probe_name', 'brain_region'], dtype='object')

ccf_table_df.sample(n=5)

import matplotlib.pyplot as plt

from iblatlas.regions import BrainRegions

COSMOS_FULL_NAMES = {
    "Isocortex": "Isocortex",
    "OLF": "Olfactory areas",
    "HPF": "Hippocampal formation",
    "CTXsp": "Cortical subplate",
    "CNU": "Cerebral nuclei",
    "TH": "Thalamus",
    "HY": "Hypothalamus",
    "MB": "Midbrain",
    "HB": "Hindbrain",
    "CB": "Cerebellum",
    "root": "Not in Cosmos",
    "void": "Outside brain",
}


def get_cosmos_color(cosmos_acronym, brain_regions=None):
    if brain_regions is None:
        brain_regions = BrainRegions()
    try:
        index = brain_regions.acronym2index(cosmos_acronym)[1][0]
        if hasattr(index, "__len__"):
            index = index[0]
        rgb = brain_regions.rgb[index]
        return (rgb[0] / 255, rgb[1] / 255, rgb[2] / 255)
    except (IndexError, KeyError):
        return (0.5, 0.5, 0.5)

nwbfile = nwbfile_raw

# %% Get anatomical localization table
anatomical_tables = localization.anatomical_coordinates_tables
ibl_bregma_table_df = anatomical_tables["ElectrodesIBLBregma"].to_dataframe()

# %% Get all unique probe names
probe_names = sorted(ibl_bregma_table_df["probe_name"].unique())
n_probes = len(probe_names)

print(f"Found {n_probes} probe(s): {probe_names}")

# %% Create figure - 2 columns per probe (contacts + regions), side by side
fig, axes = plt.subplots(1, n_probes * 2, figsize=(5 * n_probes, 12), sharey=True)

# Handle single probe case (axes won't be a list)
if n_probes == 1:
    axes = [axes[0], axes[1]]

# Track all unique regions across all probes for the legend
all_unique_regions = []

# %% Plot each probe
for probe_index, probe_name in enumerate(probe_names):
    # Get axes for this probe
    ax_probe = axes[probe_index * 2]
    ax_regions = axes[probe_index * 2 + 1]

    # Filter data for this probe
    probe_table = ibl_bregma_table_df[ibl_bregma_table_df["probe_name"] == probe_name].copy()

    # Extract electrode positions from localized_entity
    probe_table["rel_x"] = [electrode["rel_x"].iloc[0] for electrode in probe_table["localized_entity"]]
    probe_table["rel_y"] = [electrode["rel_y"].iloc[0] for electrode in probe_table["localized_entity"]]

    # Sort by depth (rel_y)
    probe_table = probe_table.sort_values("rel_y").reset_index(drop=True)

    # Precompute colors
    cosmos_values = probe_table["cosmos_location"].values
    rel_x_values = probe_table["rel_x"].values
    rel_y_values = probe_table["rel_y"].values
    colors = [get_cosmos_color(c) for c in cosmos_values]

    # Track unique regions for legend
    for region in cosmos_values:
        if region not in all_unique_regions:
            all_unique_regions.append(region)

    # --- Left panel: Probe contacts ---
    contact_width = 12  # um
    contact_height = 10  # um

    for i in range(len(rel_x_values)):
        rect = plt.Rectangle(
            (rel_x_values[i] - contact_width / 2, rel_y_values[i] - contact_height / 2),
            contact_width,
            contact_height,
            facecolor=colors[i],
            edgecolor="black",
            linewidth=0.3,
        )
        ax_probe.add_patch(rect)

    # Set probe axis limits
    x_min, x_max = probe_table["rel_x"].min(), probe_table["rel_x"].max()
    y_min, y_max = probe_table["rel_y"].min(), probe_table["rel_y"].max()
    x_margin, y_margin = 50, 100

    ax_probe.set_xlim(x_min - x_margin, x_max + x_margin)
    ax_probe.set_ylim(y_min - y_margin, y_max + y_margin)
    
    # Simplify x-axis: single tick at max value
    ax_probe.set_xticks([x_max])
    ax_probe.set_xticklabels([f"{x_max:.0f}"])
    ax_probe.set_xlabel("Relative X (um)")
    
    if probe_index == 0:
        ax_probe.set_ylabel("Distance to Probe Tip (um)")
    ax_probe.set_title(f"{probe_name}")
    ax_probe.set_aspect("equal", adjustable="box")

    # --- Right panel: Region blocks ---
    blocks = []
    current_region = cosmos_values[0]
    block_start = rel_y_values[0]

    for i in range(1, len(cosmos_values)):
        if cosmos_values[i] != current_region:
            block_end = (rel_y_values[i - 1] + rel_y_values[i]) / 2
            blocks.append((current_region, block_start, block_end))
            current_region = cosmos_values[i]
            block_start = block_end

    # Add final block
    blocks.append((current_region, block_start, rel_y_values[-1] + contact_height / 2))

    # Plot region blocks
    block_width = 1.0
    for region, y_start, y_end in blocks:
        color = get_cosmos_color(region)
        height = y_end - y_start

        rect = plt.Rectangle(
            (0, y_start),
            block_width,
            height,
            facecolor=color,
            edgecolor="black",
            linewidth=0.5,
        )
        ax_regions.add_patch(rect)

        # Add label in the middle of the block
        y_center = (y_start + y_end) / 2
        display_name = COSMOS_FULL_NAMES.get(region, region)
        ax_regions.text(block_width + 0.1, y_center, display_name, va="center", ha="left", fontsize=9, fontweight="bold")

    ax_regions.set_xlim(-0.2, 2.0)
    ax_regions.set_ylim(y_min - y_margin, y_max + y_margin)
    ax_regions.set_xticks([])
    ax_regions.set_title("Cosmos Regions")
    ax_regions.axis("off")

# %% Add legend and title
legend_handles = [
    plt.Rectangle((0, 0), 1, 1, facecolor=get_cosmos_color(r), edgecolor="black", linewidth=0.5)
    for r in all_unique_regions
]
legend_labels = [COSMOS_FULL_NAMES.get(r, r) for r in all_unique_regions]


plt.suptitle(f"Session: {nwbfile.session_id}", fontsize=14, fontweight="bold")
plt.tight_layout(rect=[0, 0.08, 1, 0.96])

plt.show()

Found 2 probe(s): ['Probe00', 'Probe01']

# Workaround: brainglobe-atlasapi 2.3.1 fetches its atlas-version manifest
# from a Colab-unfriendly S3 endpoint that returns 403 for some Colab IPs.
# Pre-seed the cache so BrainGlobeAtlas() below can find a version manifest
# even when the live fetch fails. See https://github.com/brainglobe/brainglobe-atlasapi/issues/738
from pathlib import Path
import urllib.request

_bg_dir = Path.home() / ".brainglobe"
_manifest_url = (
    "https://brainglobe.s3.us-west-2.amazonaws.com/atlas/atlases/last_versions.conf"
)
try:
    _manifest = urllib.request.urlopen(_manifest_url, timeout=20).read().decode()
except Exception as _e:  # noqa: BLE001
    print(f"Couldn't reach brainglobe manifest ({_e}); writing a minimal fallback.")
    _manifest = "[atlases]\nallen_mouse_25um = 3.0\n"

# Write to both the V1 and V2 cache paths so whichever brainglobe-atlasapi
# code path reads from disk finds the file.
for _p in (_bg_dir / "last_versions.conf",
           _bg_dir / "brainglobe-atlasapi" / "atlases" / "last_versions.conf"):
    _p.parent.mkdir(parents=True, exist_ok=True)
    if not _p.exists():
        _p.write_text(_manifest)

import matplotlib.pyplot as plt
import numpy as np
from brainglobe_atlasapi import BrainGlobeAtlas

# %% Get CCFv3 coordinates for all probes
anatomical_tables = localization.anatomical_coordinates_tables
ccf_table_df = anatomical_tables["ElectrodesCCFv3"].to_dataframe()

# Get all unique probe names
probe_names = sorted(ccf_table_df["probe_name"].unique())
n_probes = len(probe_names)

# Define colors for each probe
probe_colors = ["red", "blue", "green", "orange", "purple"][:n_probes]

# %% Load atlas
atlas = BrainGlobeAtlas("allen_mouse_25um")

# %% Collect probe data - compute slice positions per probe (mean coordinate approach)
# See documentation/probe_slice_visualization_guide.md for heuristics
probe_data = {}

for probe_name in probe_names:
    probe_table = ccf_table_df[ccf_table_df["probe_name"] == probe_name].copy()
    bg_x = probe_table["x"].values  # Anterior (0) to Posterior (13200)
    bg_y = probe_table["y"].values  # Dorsal (0) to Ventral (8000)
    bg_z = probe_table["z"].values  # Left (0) to Right (11400)
    
    # Calculate slice positions for this specific probe (mean coordinate)
    coronal_position = int(np.mean(bg_x))
    sagittal_position = int(np.mean(bg_z))
    horizontal_position = int(np.mean(bg_y))
    
    # Convert to slice indices
    coronal_index = int(coronal_position / atlas.resolution[0])
    sagittal_index = int(sagittal_position / atlas.resolution[2])
    horizontal_index = int(horizontal_position / atlas.resolution[1])
    
    probe_data[probe_name] = {
        "x_px": bg_x / atlas.resolution[0],
        "y_px": bg_y / atlas.resolution[1],
        "z_px": bg_z / atlas.resolution[2],
        "coronal_position": coronal_position,
        "sagittal_position": sagittal_position,
        "horizontal_position": horizontal_position,
        "coronal_index": coronal_index,
        "sagittal_index": sagittal_index,
        "horizontal_index": horizontal_index,
    }

# %% Create Figure: Grid with n_probes rows x 3 columns (one row per probe)
fig, axes = plt.subplots(n_probes, 3, figsize=(20, 7 * n_probes))

# Handle single probe case (axes won't be 2D)
if n_probes == 1:
    axes = axes.reshape(1, -1)

for row, probe_name in enumerate(probe_names):
    data = probe_data[probe_name]
    color = probe_colors[row]
    
    # Get slices centered on this probe's location
    reference_coronal = atlas.reference[data["coronal_index"], :, :]
    reference_sagittal = atlas.reference[:, :, data["sagittal_index"]]
    reference_horizontal = atlas.reference[:, data["horizontal_index"], :]
    
    # Coronal view
    ax = axes[row, 0]
    ax.imshow(reference_coronal, cmap="gray", aspect="equal")
    ax.scatter(data["z_px"], data["y_px"], c=color, s=10, alpha=0.8)
    ax.set_xlabel("Left - Right (pixels)")
    ax.set_ylabel("Dorsal - Ventral (pixels)")
    ax.set_title(f"{probe_name} - Coronal (AP = {data['coronal_position']} um)")
    
    # Sagittal view (transposed for correct orientation)
    ax = axes[row, 1]
    ax.imshow(reference_sagittal.T, cmap="gray", aspect="equal")
    ax.scatter(data["x_px"], data["y_px"], c=color, s=10, alpha=0.8)
    ax.set_xlabel("Anterior - Posterior (pixels)")
    ax.set_ylabel("Dorsal - Ventral (pixels)")
    ax.set_title(f"{probe_name} - Sagittal (ML = {data['sagittal_position']} um)")
    
    # Horizontal view
    ax = axes[row, 2]
    ax.imshow(reference_horizontal, cmap="gray", aspect="equal")
    ax.scatter(data["z_px"], data["x_px"], c=color, s=10, alpha=0.8)
    ax.set_xlabel("Left - Right (pixels)")
    ax.set_ylabel("Anterior - Posterior (pixels)")
    ax.set_title(f"{probe_name} - Horizontal (DV = {data['horizontal_position']} um)")

fig.suptitle(f"Session: {nwbfile_raw.session_id} - Probe Trajectories (Reference)", fontsize=14, fontweight="bold")
plt.tight_layout()
plt.show()

Could not fetch the latest atlas versions: 403
Using the last cached version from /home/runner/.brainglobe/last_versions.conf

Could not fetch the latest atlas versions: 403
Using the last cached version from /home/runner/.brainglobe/last_versions.conf

Could not fetch the latest atlas versions: 403
Using the last cached version from /home/runner/.brainglobe/last_versions.conf

Could not fetch the latest atlas versions: 403
Using the last cached version from /home/runner/.brainglobe/last_versions.conf

---------------------------------------------------------------------------
ReadError                                 Traceback (most recent call last)
/opt/hostedtoolcache/Python/3.13.15/x64/lib/python3.13/site-packages/brainglobe_atlasapi/bg_atlas.py in download_extract_file(self)
    203         try:
--> 204             tar = tarfile.open(destination_path)
    205             tar.extractall(path=self.brainglobe_dir)

/opt/hostedtoolcache/Python/3.13.15/x64/lib/python3.13/tarfile.py in open(cls, name, mode, fileobj, bufsize, **kwargs)
   1943             error_msgs_summary = '\n'.join(error_msgs)
-> 1944             raise ReadError(f"file could not be opened successfully:\n{error_msgs_summary}")
   1945 

ReadError: file could not be opened successfully:
- method gz: ReadError('not a gzip file')
- method bz2: ReadError('not a bzip2 file')
- method xz: ReadError('not an lzma file')
- method tar: ReadError('truncated header')

The above exception was the direct cause of the following exception:

ReadError                                 Traceback (most recent call last)
/tmp/ipykernel_2421/3983204879.py in <cell line: 0>()
     40 
     41 # %% Load atlas
---> 42 atlas = BrainGlobeAtlas("allen_mouse_25um")
     43 
     44 # %% Collect probe data - compute slice positions per probe (mean coordinate approach)

/opt/hostedtoolcache/Python/3.13.15/x64/lib/python3.13/site-packages/brainglobe_atlasapi/bg_atlas.py in __init__(self, atlas_name, brainglobe_dir, interm_download_dir, check_latest, config_dir, fn_update)
     97                 raise ValueError(f"{atlas_name} is not a valid atlas name!")
     98             else:
---> 99                 self.download_extract_file()
    100 
    101         # Instantiate after eventual download:

/opt/hostedtoolcache/Python/3.13.15/x64/lib/python3.13/site-packages/brainglobe_atlasapi/bg_atlas.py in download_extract_file(self)
    208             # Delete corrupted file so we can try again later
    209             destination_path.unlink(missing_ok=True)
--> 210             raise tarfile.ReadError(
    211                 f"Atlas download was interrupted or corrupted. "
    212                 f"Please try again. Details: {e}"

ReadError: Atlas download was interrupted or corrupted. Please try again. Details: file could not be opened successfully:
- method gz: ReadError('not a gzip file')
- method bz2: ReadError('not a bzip2 file')
- method xz: ReadError('not an lzma file')
- method tar: ReadError('truncated header')

from pathlib import Path

from pynwb.image import ImageSeries

# Videos in NWB are stored as ImageSeries with external_file containing relative paths
# We resolve them to full DANDI asset paths and get S3 URLs for streaming

nwb_parent = Path(raw_asset.path).parent
video_s3_urls = {}

for name, obj in nwbfile_raw.acquisition.items():
    # Check for ImageSeries with external_file (videos stored externally)
    if isinstance(obj, ImageSeries) and obj.external_file is not None:
        # Get the relative path and resolve to full DANDI path
        relative_path = obj.external_file[0].lstrip("./")
        full_path = str(nwb_parent / relative_path)
        
        # Fetch the video asset and get its S3 URL
        video_asset = dandiset.get_asset_by_path(full_path)
        if video_asset is not None:
            video_s3_urls[name] = video_asset.get_content_url(
                follow_redirects=1, strip_query=True
            )

video_s3_urls

{'VideoBodyCamera': 'https://dandiarchive.s3.amazonaws.com/blobs/710/ab6/710ab6a6-0092-463d-80d7-c7c19966a380',
 'VideoLeftCamera': 'https://dandiarchive.s3.amazonaws.com/blobs/8bd/94f/8bd94ffa-07c3-4889-9ee8-adf26f1761b2',
 'VideoRightCamera': 'https://dandiarchive.s3.amazonaws.com/blobs/c9f/c05/c9fc050a-a7f0-4d97-bd87-a076eb483e8a'}

# Stream a video directly in the notebook using IPython HTML display
# This works but might have limitations in some notebook environments
from IPython.display import HTML

camera_to_display = "VideoLeftCamera"

if camera_to_display in video_s3_urls:
    video_url = video_s3_urls[camera_to_display]
    display(HTML(f"""
    <video width="640" height="480" controls>
        <source src="{video_url}" type="video/mp4">
        Your browser does not support the video tag.
    </video>
    <p><b>{camera_to_display}</b></p>
    """))
else:
    print(f"{camera_to_display} not available for this session")

from nwb_video_widgets import NWBDANDIVideoPlayer

video_grid=[["VideoLeftCamera", "VideoBodyCamera", "VideoRightCamera"]]
NWBDANDIVideoPlayer(asset=raw_asset, nwbfile=nwbfile_raw, video_grid = video_grid)

s3_url = processed_asset.get_content_url(follow_redirects=1, strip_query=False)
file_system = remfile.File(s3_url)
file = h5py.File(file_system, mode="r")

io = NWBHDF5IO(file=file)
nwbfile_processed = io.read()

/opt/hostedtoolcache/Python/3.13.15/x64/lib/python3.13/site-packages/hdmf/spec/namespace.py:620: UserWarning: Ignoring the following cached namespace(s) because another version is already loaded:
hdmf-common - cached version: 1.9.0, loaded version: 1.8.0
hdmf-experimental - cached version: 0.6.0, loaded version: 0.5.0
Please update to the latest package versions.
  self.warn_for_ignored_namespaces(ignored_namespaces)

nwbfile_processed

session_start_time = nwbfile_processed.session_start_time
eid = nwbfile_processed.session_id
number_of_trials = len(nwbfile_processed.trials)
number_of_units = len(nwbfile_processed.units)
lab = nwbfile_processed.lab

print(f"Session ID: {eid}")
print(f"Session Start Time: {session_start_time}")
print(f"Number of Trials: {number_of_trials}")
print(f"Number of Units: {number_of_units}")
print(f"Lab: {lab}")

Session ID: 78b4fff5-c5ec-44d9-b5f9-d59493063f00
Session Start Time: 2022-07-04 14:30:15.049000+01:00
Number of Trials: 892
Number of Units: 2245
Lab: cortexlab

subject = nwbfile_processed.subject
subject

nwbfile.epochs

epochs_df = nwbfile.epochs.to_dataframe()
epochs_df

units = nwbfile_processed.units
units

units_df = units.to_dataframe()
units_df.columns

Index(['unit_name', 'electrodes', 'min_spike_amplitude_uV', 'max_electrode',
       'spike_amplitudes_uV', 'spike_times', 'isi_violations_ratio',
       'spike_amplitude_std_dB', 'median_spike_amplitude_uV',
       'cumulative_drift_um_per_hour', 'presence_ratio_std', 'firing_rate',
       'presence_ratio', 'peak_to_trough_duration_ms',
       'spike_distances_from_probe_tip_um', 'kilosort2_label', 'waveform_mean',
       'probe_name', 'sliding_rp_violation', 'rp_violation',
       'missed_spikes_estimate', 'cluster_uuid', 'noise_cutoff',
       'distance_from_probe_tip_um', 'spike_count', 'max_spike_amplitude_uV',
       'ibl_quality_score'],
      dtype='object')

probe_names = sorted(units_df["probe_name"].unique())
print(f"Available probes: {probe_names}")

# Select the first probe for detailed analysis
selected_probe = probe_names[0]
units_probe_df = units_df[units_df["probe_name"] == selected_probe]
print(f"Analyzing {selected_probe}: {len(units_probe_df)} units")

columns = ["unit_name", "cluster_uuid", "spike_count", "firing_rate", "kilosort2_label", "ibl_quality_score"]
units_probe_df[columns].sample(n=5)

Available probes: ['Probe00', 'Probe01']
Analyzing Probe00: 684 units

columns = ["unit_name", "median_spike_amplitude_uV", "distance_from_probe_tip_um"]
units_probe_df[columns].sample(n=5)

task_epoch = epochs_df[epochs_df['protocol_type'] == 'task'].iloc[0]
passive_epoch = epochs_df[epochs_df['protocol_type'] == 'passive'].iloc[0]

# Calculate durations for labels
task_duration = task_epoch['stop_time'] - task_epoch['start_time']
passive_duration = passive_epoch['stop_time'] - passive_epoch['start_time']

# Filter to good units only (IBL quality score == 1)
df = units_probe_df[units_probe_df['ibl_quality_score'] == 1].copy()
print(f"Number of good units on {probe_name}: {len(units_probe_df[units_probe_df['ibl_quality_score'] == 1])}")

# Sort units by depth
units_probe = df.sort_values('distance_from_probe_tip_um')

# Create figure
fig, ax = plt.subplots(figsize=(14, 8))

# Add epoch background colors with duration in labels
ax.axvspan(task_epoch['start_time'], task_epoch['stop_time'],
           color='lightblue', alpha=0.3, label=f'Task ({task_duration/60:.1f} min)')
ax.axvspan(passive_epoch['start_time'], passive_epoch['stop_time'],
           color='lightyellow', alpha=0.3, label=f'Passive ({passive_duration/60:.1f} min)')

# Plot spikes for each unit at its depth
for idx, (unit_id, unit) in enumerate(units_probe.iterrows()):
    spike_times = unit['spike_times']
    depth = unit['distance_from_probe_tip_um']

    # Plot all spikes
    ax.scatter(spike_times, np.full(len(spike_times), depth),
               c='black', s=1.0, marker='|', linewidths=0.5, alpha=0.5)

# Add epoch boundary line
ax.axvline(task_epoch['stop_time'], color='red', linestyle='--', linewidth=1.5, label='Task/Passive boundary')

# Labels and styling
ax.set_xlabel('Time (s)')
ax.set_ylabel('Distance From Probe Tip (um)')
ax.set_title(f'Spike Raster by Depth - {probe_name} (good units only)',
             fontsize=12, fontweight='bold')
ax.set_xlim(0, passive_epoch['stop_time'])
ax.set_ylim(0, units_probe['distance_from_probe_tip_um'].max() + 100)

# Keep y-axis normal (depth increases upward, deeper units at top)

# Add legend outside the plot
ax.legend(loc='upper left', bbox_to_anchor=(1.01, 1), fontsize=9)

plt.tight_layout()
plt.show()

Number of good units on Probe01: 58

import matplotlib.pyplot as plt

fig, axes = plt.subplots(1, 2, figsize=(12, 9), sharey=True)

# Left plot - colored by kilosort2_label
df = units_probe_df
value_to_color = {"mua": "red", "good": "green"}

for label, color in value_to_color.items():
    subset = df[df["kilosort2_label"] == label]
    axes[0].scatter(
        subset["median_spike_amplitude_uV"],
        subset["distance_from_probe_tip_um"],
        c=color,
        label=label,
        s=100,
    )
axes[0].set_xlim([0, 800])
axes[0].set_xlabel("median_spike_amplitude_uV")
axes[0].set_ylabel("mean_relative_depth_um")
axes[0].set_title("By Kilosort2 Label")
axes[0].legend()

# Right plot - colored by firing_rate
df = units_probe_df[units_probe_df["ibl_quality_score"] == 1]

df.plot.scatter(
    x="median_spike_amplitude_uV", 
    y="distance_from_probe_tip_um",
    c="firing_rate",
    colormap="inferno",
    s=100,
    ax=axes[1]
)
axes[1].set_xlim([0, 800])
axes[1].set_xlabel("median_spike_amplitude_uV")
axes[1].set_title("By Firing Rate (Quality=1)")

plt.tight_layout()
plt.show()

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Patch

# COSMOS_FULL_NAMES and get_cosmos_color defined above


def get_cosmos_region_data(probe_table):
    """Extract COSMOS region blocks from probe table for visualization."""
    probe_table = probe_table.sort_values("rel_y").reset_index(drop=True)
    cosmos_values = probe_table["cosmos_location"].values
    rel_y_values = probe_table["rel_y"].values

    blocks = []
    current_region = cosmos_values[0]
    block_start = rel_y_values[0]

    for i in range(1, len(cosmos_values)):
        if cosmos_values[i] != current_region:
            block_end = (rel_y_values[i - 1] + rel_y_values[i]) / 2
            blocks.append((current_region, block_start, block_end))
            current_region = cosmos_values[i]
            block_start = block_end

    blocks.append((current_region, block_start, rel_y_values[-1] + 5))

    unique_regions = []
    for region, _, _ in blocks:
        if region not in unique_regions:
            unique_regions.append(region)

    return blocks, unique_regions


def plot_probe_regions(probe_table, ax, unit_depths=None, unit_ids=None):
    """Plot probe with COSMOS region strip and electrode contacts."""
    probe_table = probe_table.sort_values("rel_y").reset_index(drop=True)
    rel_x_values = probe_table["rel_x"].values
    rel_y_values = probe_table["rel_y"].values

    x_min, x_max = rel_x_values.min(), rel_x_values.max()
    y_min, y_max = rel_y_values.min(), rel_y_values.max()

    blocks, unique_regions = get_cosmos_region_data(probe_table)
    strip_width = x_max - x_min + 30
    strip_x = x_min - 15

    for region, y_start, y_end in blocks:
        color = get_cosmos_color(region)
        rect = plt.Rectangle(
            (strip_x, y_start), strip_width, y_end - y_start, facecolor=color, edgecolor="none", alpha=0.5
        )
        ax.add_patch(rect)

    for i in range(len(rel_x_values)):
        rect = plt.Rectangle(
            (rel_x_values[i] - 6, rel_y_values[i] - 5), 12, 10, facecolor="white", edgecolor="black", linewidth=0.3
        )
        ax.add_patch(rect)

    if unit_depths is not None and unit_ids is not None:
        for i, depth in enumerate(unit_depths):
            if i < len(unit_ids):
                ax.text(x_max + 25, depth, str(unit_ids[i]), va="center", ha="left", fontsize=7, color="black", fontweight="bold")

    ax.set_xlim(x_min - 50, x_max + 50)
    ax.set_ylim(y_min - 100, y_max + 100)
    ax.set_xticks([])
    ax.set_ylabel("Depth (um)")
    ax.set_aspect("equal", adjustable="box")
    for spine in ax.spines.values():
        spine.set_visible(False)

    legend_handles = [
        Patch(facecolor=get_cosmos_color(r), edgecolor="black", linewidth=0.5, label=COSMOS_FULL_NAMES.get(r, r), alpha=0.5)
        for r in unique_regions
    ]
    return ax, legend_handles


def plot_waveform_wiggle(waveform_uV, electrode_depths, ax, fs=30000.0, gain=0.4, channel_spacing=0.8):
    """Plot waveform as IBL-style double wiggle plot.

    Parameters
    ----------
    waveform_uV : array
        Waveform data in microvolts (num_samples, num_channels).
    electrode_depths : array
        Depth of each electrode channel.
    ax : matplotlib axis
    fs : float
        Sampling rate in Hz.
    gain : float
        Scaling factor for waveform amplitude.
    channel_spacing : float
        Vertical spacing between channels.
    """
    num_samples, num_channels = waveform_uV.shape

    # Handle mismatch between waveform channels and electrode depths
    n_electrodes = len(electrode_depths)
    if num_channels != n_electrodes:
        # Use only channels that have electrode depths
        num_channels = min(num_channels, n_electrodes)
        waveform_uV = waveform_uV[:, :num_channels]
        electrode_depths = electrode_depths[:num_channels]

    center_sample = num_samples // 2
    time_ms = (np.arange(num_samples) - center_sample) / fs * 1000

    sort_order = np.argsort(electrode_depths)
    wf_sorted = waveform_uV[:, sort_order]
    depths_sorted = electrode_depths[sort_order]

    wf_max = np.abs(wf_sorted).max()
    wf_norm = wf_sorted / wf_max * gain if wf_max > 0 else wf_sorted

    for ch in range(num_channels):
        offset = ch * channel_spacing
        trace = wf_norm[:, ch]
        ax.plot(time_ms, trace + offset, color="black", linewidth=0.5)
        ax.fill_between(time_ms, offset, trace + offset, where=(trace > 0), color="orange", alpha=0.8, linewidth=0)
        ax.fill_between(time_ms, offset, trace + offset, where=(trace < 0), color="slateblue", alpha=0.8, linewidth=0)

    ax.set_xlabel("Time (ms)")
    depth_ticks = np.linspace(depths_sorted.min(), depths_sorted.max(), 5)
    channel_positions = np.linspace(0, (num_channels - 1) * channel_spacing, 5)
    ax.set_yticks(channel_positions)
    ax.set_yticklabels([f"{d:.0f}" for d in depth_ticks])
    ax.set_ylabel("Distance from tip (um)")


def get_max_channel_index(unit):
    """Get the index (0-31) of the max electrode within the unit's electrodes list."""
    max_electrode_idx = unit["max_electrode"].index[0]
    electrodes_indices = unit["electrodes"].index.tolist()
    if max_electrode_idx in electrodes_indices:
        return electrodes_indices.index(max_electrode_idx)
    waveform = unit["waveform_mean"]
    peak_to_peak = waveform.max(axis=0) - waveform.min(axis=0)
    return np.argmax(peak_to_peak)


def plot_max_channel_waveform(waveform_uV, ax, max_ch_idx=None, fs=30000.0):
    """Plot waveform on the max amplitude channel.

    Parameters
    ----------
    waveform_uV : array
        Waveform data in microvolts (num_samples, num_channels).
    ax : matplotlib axis
    max_ch_idx : int, optional
        Index of max channel in the waveform array. If None, computed from waveform.
    fs : float
        Sampling rate in Hz.
    """
    num_samples = waveform_uV.shape[0]
    center_sample = num_samples // 2
    time_ms = (np.arange(num_samples) - center_sample) / fs * 1000

    if max_ch_idx is None:
        peak_to_peak = waveform_uV.max(axis=0) - waveform_uV.min(axis=0)
        max_ch_idx = np.argmax(peak_to_peak)

    trace = waveform_uV[:, max_ch_idx]

    ax.plot(time_ms, trace, color="black", linewidth=1.0)
    ax.fill_between(time_ms, 0, trace, where=(trace > 0), color="orange", alpha=0.8, linewidth=0)
    ax.fill_between(time_ms, 0, trace, where=(trace < 0), color="slateblue", alpha=0.8, linewidth=0)
    ax.axhline(0, color="gray", linestyle="--", linewidth=0.5)
    ax.set_xlabel("Time (ms)")
    ax.set_ylabel("Amplitude (uV)")


def plot_waveform_heatmap(waveform_uV, electrode_depths, ax, fs=30000.0, vmax=None, add_colorbar=False):
    """Plot waveform as a heatmap (time x channel).

    Parameters
    ----------
    waveform_uV : array
        Waveform data in microvolts (num_samples, num_channels).
    electrode_depths : array
        Depth of each electrode channel.
    ax : matplotlib axis
    fs : float
        Sampling rate in Hz.
    vmax : float, optional
        Max value for colormap.
    add_colorbar : bool
        Whether to add a colorbar.
    """
    num_samples, num_channels = waveform_uV.shape

    # Handle mismatch between waveform channels and electrode depths
    n_electrodes = len(electrode_depths)
    if num_channels != n_electrodes:
        num_channels = min(num_channels, n_electrodes)
        waveform_uV = waveform_uV[:, :num_channels]
        electrode_depths = electrode_depths[:num_channels]

    center_sample = num_samples // 2
    time_ms = (np.arange(num_samples) - center_sample) / fs * 1000

    sort_order = np.argsort(electrode_depths)
    wf_sorted = waveform_uV[:, sort_order]
    depths_sorted = electrode_depths[sort_order]

    if vmax is None:
        vmax = np.abs(wf_sorted).max()

    im = ax.imshow(
        wf_sorted.T,
        aspect="auto",
        origin="lower",
        extent=[time_ms[0], time_ms[-1], 0, num_channels],
        cmap="RdBu_r",
        vmin=-vmax,
        vmax=vmax,
    )
    ax.set_xlabel("Time (ms)")

    depth_ticks = np.linspace(depths_sorted.min(), depths_sorted.max(), 5)
    channel_positions = np.linspace(0, num_channels, 5)
    ax.set_yticks(channel_positions)
    ax.set_yticklabels([f"{d:.0f}" for d in depth_ticks])
    ax.set_ylabel("Distance from tip (um)")

    if add_colorbar:
        plt.colorbar(im, ax=ax, label="Amplitude (uV)")
    return im

# %% Select good units and create combined waveform visualization
# Get localization data for probe visualization
localization = nwbfile_processed.lab_meta_data.get("localization")
anatomical_tables = localization.anatomical_coordinates_tables
ibl_bregma_table_df = anatomical_tables["ElectrodesIBLBregma"].to_dataframe()

# Extract electrode positions from localized_entity
ibl_bregma_table_df["rel_x"] = [electrode["rel_x"].iloc[0] for electrode in ibl_bregma_table_df["localized_entity"]]
ibl_bregma_table_df["rel_y"] = [electrode["rel_y"].iloc[0] for electrode in ibl_bregma_table_df["localized_entity"]]

# Select good units from the first probe that has them
probe_names = sorted(units_df["probe_name"].unique())
probe_name = None
good_units = None
for candidate in probe_names:
    candidate_units = units_df[(units_df["probe_name"] == candidate) & (units_df["ibl_quality_score"] == 1.0)].copy()
    if len(candidate_units) > 0:
        probe_name = candidate
        good_units = candidate_units
        break

if good_units is None or len(good_units) == 0:
    raise ValueError(f"No good units found in any probe: {probe_names}")

good_units = good_units.sort_values("distance_from_probe_tip_um")
probe_table = ibl_bregma_table_df[ibl_bregma_table_df["probe_name"] == probe_name].copy()

# Filter out units with any NaN waveforms (some units have partial NaN data)
valid_waveform_mask = [not np.isnan(unit["waveform_mean"]).any() for _, unit in good_units.iterrows()]
good_units = good_units[valid_waveform_mask]

print(f"Selected {len(good_units)} good units with valid waveforms from {probe_name}")

# Select 3 units evenly spaced across depth
depths = good_units["distance_from_probe_tip_um"].values
target_depths = np.linspace(depths.min(), depths.max(), 3)
selected_indices = []
available_mask = np.ones(len(good_units), dtype=bool)

for target in target_depths:
    available_depths = depths.copy()
    available_depths[~available_mask] = np.inf
    closest_idx = np.argmin(np.abs(available_depths - target))
    selected_indices.append(closest_idx)
    available_mask[closest_idx] = False

selected_indices = sorted(selected_indices, key=lambda i: depths[i])
example_units = good_units.iloc[selected_indices]
example_depths = example_units["distance_from_probe_tip_um"].values.tolist()
example_unit_ids = example_units.index.tolist()

# Convert waveforms from volts to microvolts for visualization
# NWB schema has waveform_mean unit fixed to 'volts', so we always convert to uV for display
conversion_factor = 1e6
print("Converting waveforms from volts to microvolts for display (factor: 1e6)")

# Compute shared vmax for heatmaps (in microvolts)
shared_vmax = max(np.abs(unit["waveform_mean"] * conversion_factor).max() for _, unit in example_units.iterrows())

# Create combined figure: Probe | 3x3 waveforms | colorbar
fig = plt.figure(figsize=(16, 12))
gs = fig.add_gridspec(3, 5, width_ratios=[1.0, 1.2, 1.2, 1.2, 0.08], wspace=0.35, hspace=0.4)

# Probe visualization (spans all 3 rows)
ax_probe = fig.add_subplot(gs[:, 0])
_, legend_handles = plot_probe_regions(probe_table, ax_probe, unit_depths=example_depths, unit_ids=example_unit_ids)
ax_probe.set_title(f"{probe_name}\nElectrode Contacts", fontsize=10)

if legend_handles:
    ax_probe.legend(handles=legend_handles, loc="upper center", bbox_to_anchor=(0.5, -0.02),
                    ncol=2, fontsize=6, frameon=False, title="COSMOS Regions", title_fontsize=7)

# Store last heatmap image for shared colorbar
last_heatmap_im = None

# Plot waveforms for each unit (columns 1-3)
for col_idx, (unit_id, unit) in enumerate(example_units.iterrows()):
    # Convert waveform to microvolts for visualization
    waveform_uV = unit["waveform_mean"] * conversion_factor
    depth = unit["distance_from_probe_tip_um"]
    region = unit["max_electrode"]["location"].values[0]
    electrode_depths = unit["electrodes"]["rel_y"].values
    
    # Get max channel index from the unit's electrode mapping
    max_ch_idx = get_max_channel_index(unit)
    max_electrode_depth = unit["max_electrode"]["rel_y"].values[0]

    # Row 0: Wiggle plot
    ax_wiggle = fig.add_subplot(gs[0, col_idx + 1])
    plot_waveform_wiggle(waveform_uV, electrode_depths, ax_wiggle, gain=0.4, channel_spacing=0.8)
    ax_wiggle.set_title(f"Unit {unit_id}\n{region}", fontsize=8)

    # Row 1: Max channel waveform (use max_electrode depth in title)
    ax_max = fig.add_subplot(gs[1, col_idx + 1])
    plot_max_channel_waveform(waveform_uV, ax_max, max_ch_idx=max_ch_idx)
    ax_max.set_title(f"Max electrode depth: {max_electrode_depth:.0f} um", fontsize=8)

    # Row 2: Heatmap
    ax_heatmap = fig.add_subplot(gs[2, col_idx + 1])
    last_heatmap_im = plot_waveform_heatmap(waveform_uV, electrode_depths, ax_heatmap, vmax=shared_vmax)

# Add shared colorbar for heatmaps
cbar_ax = fig.add_subplot(gs[2, 4])
cbar = fig.colorbar(last_heatmap_im, cax=cbar_ax)
cbar.set_label("Amplitude (uV)", fontsize=9)

fig.suptitle(f"Waveform Visualization - {probe_name}\nSession: {nwbfile_processed.session_id}",
             fontsize=12, fontweight="bold", y=1.02)
plt.tight_layout()
plt.show()

Selected 51 good units with valid waveforms from Probe00
Converting waveforms from volts to microvolts for display (factor: 1e6)

/tmp/ipykernel_2421/3968331865.py:109: UserWarning: This figure includes Axes that are not compatible with tight_layout, so results might be incorrect.
  plt.tight_layout()

trials = nwbfile_processed.trials
trials

trials_df = trials.to_dataframe()
trials_df

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

# %% Compute signed contrast for y-position
contrast = trials_df["gabor_stimulus_contrast"].values
side = trials_df["gabor_stimulus_side"].values
signed_contrast = np.where(side == "left", contrast, -contrast)
trials_df["signed_contrast"] = signed_contrast

# Add jitter for overlapping points
np.random.seed(42)
trials_df["y_jitter"] = signed_contrast + np.random.uniform(-3, 3, len(trials_df))

# %% Define block background colors
block_bg_colors = {
    "unbiased": "#E0E0E0",
    "left_block": "#B3D9FF",
    "right_block": "#FFD9B3",
}
block_labels = {
    "unbiased": "Unbiased (p_left=0.5)",
    "left_block": "Left block (p_left=0.8)",
    "right_block": "Right block (p_left=0.2)",
}

# %% Create figure
fig, ax = plt.subplots(figsize=(14, 6))

# Plot block backgrounds
block_changes = trials_df["block_index"].diff().fillna(1) != 0
block_starts = trials_df.index[block_changes].tolist()
block_starts.append(len(trials_df))

added_labels = set()
for i in range(len(block_starts) - 1):
    start_index = block_starts[i]
    end_index = block_starts[i + 1]
    block_type = trials_df.loc[start_index, "block_type"]
    label = block_labels.get(block_type) if block_type not in added_labels else None
    ax.axvspan(
        start_index,
        end_index,
        color=block_bg_colors.get(block_type, "white"),
        alpha=0.5,
        zorder=0,
        label=label,
    )
    added_labels.add(block_type)

# Plot correct trials (blue circles)
correct = trials_df[trials_df["is_mouse_rewarded"]]
ax.scatter(
    correct.index,
    correct["y_jitter"],
    marker="o",
    s=20,
    facecolors="none",
    edgecolors="blue",
    linewidths=1,
    label="Correct",
    zorder=2,
)

# Plot incorrect trials (red X)
incorrect = trials_df[~trials_df["is_mouse_rewarded"]]
ax.scatter(
    incorrect.index,
    incorrect["y_jitter"],
    color="red",
    marker="x",
    s=20,
    linewidths=1,
    label="Incorrect",
    zorder=2,
)

# Formatting
ax.set_xlabel("Trial number")
ax.set_ylabel("Signed contrast (%)")
ax.set_xlim(-5, len(trials_df) + 5)
ax.axhline(0, color="gray", linestyle="--", linewidth=0.5, alpha=0.5)

# Session info
subject = nwbfile.subject.subject_id if nwbfile.subject else "Unknown"
session_date = nwbfile.session_start_time.strftime("%Y-%m-%d")
n_correct = len(correct)
performance = n_correct / len(trials_df) * 100

ax.set_title(f"{subject} - {session_date} ({len(trials_df)} trials, {performance:.1f}% correct)")
ax.legend(loc="upper left", bbox_to_anchor=(1.01, 1), fontsize=9)

plt.tight_layout(rect=[0, 0, 0.85, 1])
plt.show()

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import psychofit as psy

# %% Compute signed contrast
# Convention: positive = left stimulus, negative = right stimulus
# This ensures P(choose right) increases with signed contrast
contrast = trials_df["gabor_stimulus_contrast"].values
side = trials_df["gabor_stimulus_side"].values
signed_contrast = np.where(side == "left", contrast, -contrast)

trials_df["signed_contrast"] = signed_contrast
trials_df["chose_right"] = (trials_df["mouse_wheel_choice"] == "right").astype(int)

# %% Define block types and colors (matching IBL style)
blocks = {
    0.5: {"color": "black", "label": "p_left=0.5"},
    0.2: {"color": "tab:red", "label": "p_left=0.2"},
    0.8: {"color": "tab:blue", "label": "p_left=0.8"},
}

# %% Fit and collect data for each block
x_fit = np.linspace(-100, 100, 200)
block_results = {}

for p_left, style in blocks.items():
    block_trials = trials_df[trials_df["probability_left"] == p_left]

    if len(block_trials) == 0:
        print(f"No trials for p_left={p_left}")
        continue

    # Aggregate by contrast
    grouped = (
        block_trials.groupby("signed_contrast")
        .agg(
            n_total=("chose_right", "count"),
            p_right=("chose_right", "mean"),
        )
        .reset_index()
    )

    # Fit psychometric curve
    data = np.vstack(
        [grouped["signed_contrast"].values, grouped["n_total"].values, grouped["p_right"].values]
    )

    pars, _ = psy.mle_fit_psycho(
        data,
        P_model="erf_psycho_2gammas",
        parstart=np.array([0.0, 40.0, 0.1, 0.1]),
        parmin=np.array([-50.0, 10.0, 0.0, 0.0]),
        parmax=np.array([50.0, 50.0, 0.2, 0.2]),
        nfits=10,
    )

    y_fit = psy.erf_psycho_2gammas(pars, x_fit)

    block_results[p_left] = {
        "grouped": grouped,
        "pars": pars,
        "y_fit": y_fit,
        "n_trials": len(block_trials),
    }

# %% Plot psychometric curves for all blocks
fig, ax = plt.subplots(figsize=(8, 6))

for p_left, style in blocks.items():
    if p_left not in block_results:
        continue

    result = block_results[p_left]
    grouped = result["grouped"]
    color = style["color"]
    label = style["label"]

    # Fitted curve
    ax.plot(x_fit, result["y_fit"], color=color, linewidth=2, label=f"{label} fit")

    # Data points
    ax.scatter(
        grouped["signed_contrast"],
        grouped["p_right"],
        color=color,
        s=50,
        alpha=0.8,
        label=f"{label} data",
    )

# Formatting
ax.set_xlabel("Contrasts")
ax.set_ylabel("Probability choosing right")
ax.set_ylim(-0.05, 1.05)
ax.legend(loc="upper left", fontsize=9)

plt.tight_layout()
plt.show()

from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

# %% Calculate reaction time and signed contrast
trials_df["reaction_time"] = (
    trials_df["wheel_movement_onset_time"] - trials_df["gabor_stimulus_onset_time"]
)

# Signed contrast: positive = left stimulus, negative = right stimulus
contrast = trials_df["gabor_stimulus_contrast"].values
side = trials_df["gabor_stimulus_side"].values
trials_df["signed_contrast"] = np.where(side == "left", contrast, -contrast)

# Get unique contrasts sorted
contrasts = sorted(trials_df["signed_contrast"].unique())

# %% Define block styles (matching IBL and psychometric_curve.py)
blocks = {
    0.5: {"color": "black", "label": "p_left=0.5"},
    0.2: {"color": "tab:red", "label": "p_left=0.2"},
    0.8: {"color": "tab:blue", "label": "p_left=0.8"},
}

# %% Plot reaction time by contrast for each block
fig, ax = plt.subplots(figsize=(8, 6))

for p_left, style in blocks.items():
    block_trials = trials_df[trials_df["probability_left"] == p_left]

    if len(block_trials) == 0:
        continue

    # Calculate median reaction time for each contrast (robust to outliers)
    median_rt = []
    contrast_vals = []

    for c in contrasts:
        contrast_trials = block_trials[block_trials["signed_contrast"] == c]
        if len(contrast_trials) > 0:
            median_rt.append(contrast_trials["reaction_time"].median())
            contrast_vals.append(c)

    ax.plot(
        contrast_vals,
        median_rt,
        "o-",
        color=style["color"],
        label=f"{style['label']} ({len(block_trials)} trials)",
        linewidth=1.5,
        markersize=6,
    )

# Formatting
ax.set_xlabel("Contrasts")
ax.set_ylabel("Reaction time (s)")
ax.set_xlim(-110, 110)
ax.set_ylim(0, None)
ax.axvline(0, color="gray", linestyle="--", linewidth=0.5, alpha=0.5)
ax.legend(loc="upper right", fontsize=9)

plt.tight_layout()
plt.show()

wheel_module = nwbfile_processed.processing["wheel"]
wheel_module

wheel_position = wheel_module["WheelPosition"]
wheel_position

pose_estimation_module = nwbfile_processed.processing["pose_estimation"]
pose_estimation_module

left_camera = pose_estimation_module["LeftCamera"]
left_camera

from nwb_video_widgets import NWBDANDIPoseEstimationWidget

NWBDANDIPoseEstimationWidget(
    asset=processed_asset,
    nwbfile=nwbfile_processed,
    video_asset=raw_asset,
    video_nwbfile=nwbfile_raw,

)

passive_module = nwbfile_processed.processing["passive_protocol"]
passive_module

gabor_table = passive_module["gabor_table"]
gabor_table

passive_intervals = passive_module["passive_intervals"]
passive_intervals

passive_task_replay = passive_module["passive_task_replay"]
passive_task_replay

# Trial-aligned left paw speed with correct/incorrect trial separation
import pynapple as nap

# Load with pynapple from memory
data = nap.NWBFile(nwbfile_processed)

# Get trials dataframe for metadata
trials_df = nwbfile_processed.trials.to_dataframe()

# Masks for correct/incorrect
correct_mask = trials_df['is_mouse_rewarded'].values

# Get left paw position and compute speed using pynapple's derivative
left_paw = data['LeftCamera/PoseEstimationSeriesLeftPaw']
velocity = left_paw.derivative()
speed_values = np.sqrt(velocity.values[:, 0]**2 + velocity.values[:, 1]**2)
speed = nap.Tsd(t=velocity.t, d=speed_values)

# Interpolate to regular grid (required by pynapple's compute_perievent)
rate = 1.0 / np.median(np.diff(speed.t))
n_points = int(((speed.t[-1] - speed.t[0]) * rate)) + 1
regular_t = nap.Ts(t=np.linspace(speed.t[0], speed.t[-1], n_points))
speed = speed.interpolate(regular_t)

# Create Ts objects for stimulus onset times
stim_times = nap.Ts(t=trials_df['gabor_stimulus_onset_time'].values)
correct_stim_times = nap.Ts(t=trials_df.loc[correct_mask, 'gabor_stimulus_onset_time'].values)
incorrect_stim_times = nap.Ts(t=trials_df.loc[~correct_mask, 'gabor_stimulus_onset_time'].values)

# Use pynapple's compute_perievent to align speed around stimulus onset
minmax = (-0.5, 1.0)  # 0.5s before to 1s after stimulus

perievent_all = nap.compute_perievent(speed, stim_times, window=minmax)
perievent_correct = nap.compute_perievent(speed, correct_stim_times, window=minmax)
perievent_incorrect = nap.compute_perievent(speed, incorrect_stim_times, window=minmax)

# Compute mean and SEM
correct_mean = np.nanmean(perievent_correct.values, axis=1)
correct_sem = np.nanstd(perievent_correct.values, axis=1) / np.sqrt(np.sum(~np.isnan(perievent_correct.values), axis=1))
incorrect_mean = np.nanmean(perievent_incorrect.values, axis=1)
incorrect_sem = np.nanstd(perievent_incorrect.values, axis=1) / np.sqrt(np.sum(~np.isnan(perievent_incorrect.values), axis=1))

time = perievent_correct.t

# Sort trials by outcome for raster
sorted_indices = np.argsort(correct_mask)
perievent_sorted = perievent_all.values[:, sorted_indices]

# Create figure
fig, axes = plt.subplots(2, 1, figsize=(4, 8), gridspec_kw={'height_ratios': [1, 3]}, sharex=True)

ax_avg = axes[0]
ax_avg.plot(time, correct_mean, color='blue', linewidth=1.5, label='Correct')
ax_avg.fill_between(time, correct_mean - correct_sem, correct_mean + correct_sem, color='blue', alpha=0.3)
ax_avg.plot(time, incorrect_mean, color='red', linewidth=1.5, label='Incorrect')
ax_avg.fill_between(time, incorrect_mean - incorrect_sem, incorrect_mean + incorrect_sem, color='red', alpha=0.3)
ax_avg.axvline(0, color='black', linestyle='--', linewidth=0.8)
ax_avg.set_ylabel('Speed (px/s)')
ax_avg.set_title('Left paw speed', fontsize=12, fontweight='bold', loc='center')
ax_avg.legend(loc='upper left', fontsize=8, framealpha=0.9)

ax_raster = axes[1]
n_trials = perievent_sorted.shape[1]
n_incorrect = (~correct_mask).sum()
n_correct = correct_mask.sum()

im = ax_raster.imshow(perievent_sorted.T, aspect='auto', cmap='gray_r',
                       extent=[time[0], time[-1], 0, n_trials],
                       vmin=0, vmax=np.nanpercentile(perievent_sorted, 95),
                       origin='lower')

# Side color bar
bar_width = 0.05
ax_bar = ax_raster.inset_axes([1.02, 0, bar_width, 1])
ax_bar.axhspan(0, n_incorrect / n_trials, color='red', alpha=0.8)
ax_bar.axhspan(n_incorrect / n_trials, 1, color='blue', alpha=0.8)
ax_bar.set_xlim(0, 1)
ax_bar.set_ylim(0, 1)
ax_bar.axis('off')
ax_bar.text(1.5, (n_incorrect / n_trials) / 2, 'incorrect', rotation=90, va='center', ha='left', fontsize=9, color='red', fontweight='bold')
ax_bar.text(1.5, n_incorrect / n_trials + (n_correct / n_trials) / 2, 'correct', rotation=90, va='center', ha='left', fontsize=9, color='blue', fontweight='bold')

ax_raster.axvline(0, color='black', linestyle='--', linewidth=0.8)
ax_raster.set_ylabel('Sorted Trial Number')
ax_raster.set_xlabel('T from Stim On (s)')

plt.tight_layout()
plt.show()

/opt/hostedtoolcache/Python/3.13.15/x64/lib/python3.13/site-packages/pynapple/core/time_index.py:151: UserWarning: timestamps are not sorted
  warn("timestamps are not sorted", UserWarning)
/opt/hostedtoolcache/Python/3.13.15/x64/lib/python3.13/site-packages/pynapple/core/base_class.py:56: UserWarning: 1 row(s) with NaN start/end time(s) were dropped.
  self.time_support = IntervalSet(start=self.index[0], end=self.index[-1])
/opt/hostedtoolcache/Python/3.13.15/x64/lib/python3.13/site-packages/pynapple/core/base_class.py:58: RuntimeWarning: divide by zero encountered in scalar divide
  self.rate = self.index.shape[0] / np.sum(

# Trial-aligned pupil diameter (z-scored) with correct/incorrect trial separation

# Get smoothed pupil diameter
pupil = data['LeftPupilDiameterSmoothed']

# Interpolate to regular grid (required by pynapple's compute_perievent)
rate = 1.0 / np.median(np.diff(pupil.t))
n_points = int(((pupil.t[-1] - pupil.t[0]) * rate)) + 1
regular_t = nap.Ts(t=np.linspace(pupil.t[0], pupil.t[-1], n_points))
pupil = pupil.interpolate(regular_t)

# Create Ts objects for stimulus onset times
stim_times = nap.Ts(t=trials_df['gabor_stimulus_onset_time'].values)
correct_stim_times = nap.Ts(t=trials_df.loc[correct_mask, 'gabor_stimulus_onset_time'].values)
incorrect_stim_times = nap.Ts(t=trials_df.loc[~correct_mask, 'gabor_stimulus_onset_time'].values)

minmax = (-0.5, 1.0)

perievent_all = nap.compute_perievent(pupil, stim_times, window=minmax)
perievent_correct = nap.compute_perievent(pupil, correct_stim_times, window=minmax)
perievent_incorrect = nap.compute_perievent(pupil, incorrect_stim_times, window=minmax)

# Z-score the data for visualization
all_values = perievent_all.values.flatten()
all_values = all_values[~np.isnan(all_values)]
mean_val = np.mean(all_values)
std_val = np.std(all_values)

perievent_all_z = (perievent_all.values - mean_val) / std_val
perievent_correct_z = (perievent_correct.values - mean_val) / std_val
perievent_incorrect_z = (perievent_incorrect.values - mean_val) / std_val

# Compute mean and SEM
correct_mean = np.nanmean(perievent_correct_z, axis=1)
correct_sem = np.nanstd(perievent_correct_z, axis=1) / np.sqrt(np.sum(~np.isnan(perievent_correct_z), axis=1))
incorrect_mean = np.nanmean(perievent_incorrect_z, axis=1)
incorrect_sem = np.nanstd(perievent_incorrect_z, axis=1) / np.sqrt(np.sum(~np.isnan(perievent_incorrect_z), axis=1))

time = perievent_correct.t

# Sort trials by outcome for raster
sorted_indices = np.argsort(correct_mask)
perievent_sorted = perievent_all_z[:, sorted_indices]

# Create figure
fig, axes = plt.subplots(2, 1, figsize=(4, 8), gridspec_kw={'height_ratios': [1, 3]}, sharex=True)

ax_avg = axes[0]
ax_avg.plot(time, correct_mean, color='blue', linewidth=1.5, label='Correct')
ax_avg.fill_between(time, correct_mean - correct_sem, correct_mean + correct_sem, color='blue', alpha=0.3)
ax_avg.plot(time, incorrect_mean, color='red', linewidth=1.5, label='Incorrect')
ax_avg.fill_between(time, incorrect_mean - incorrect_sem, incorrect_mean + incorrect_sem, color='red', alpha=0.3)
ax_avg.axvline(0, color='black', linestyle='--', linewidth=0.8)
ax_avg.set_ylabel('Pupil (z-score)')
ax_avg.set_title('Pupil diameter', fontsize=12, fontweight='bold', loc='center')
ax_avg.legend(loc='upper left', fontsize=8, framealpha=0.9)

ax_raster = axes[1]
n_trials = perievent_sorted.shape[1]
n_incorrect = (~correct_mask).sum()
n_correct = correct_mask.sum()

im = ax_raster.imshow(perievent_sorted.T, aspect='auto', cmap='gray_r',
                       extent=[time[0], time[-1], 0, n_trials],
                       vmin=np.nanpercentile(perievent_sorted, 5),
                       vmax=np.nanpercentile(perievent_sorted, 95),
                       origin='lower')

# Side color bar
bar_width = 0.05
ax_bar = ax_raster.inset_axes([1.02, 0, bar_width, 1])
ax_bar.axhspan(0, n_incorrect / n_trials, color='red', alpha=0.8)
ax_bar.axhspan(n_incorrect / n_trials, 1, color='blue', alpha=0.8)
ax_bar.set_xlim(0, 1)
ax_bar.set_ylim(0, 1)
ax_bar.axis('off')
ax_bar.text(1.5, (n_incorrect / n_trials) / 2, 'incorrect', rotation=90, va='center', ha='left', fontsize=9, color='red', fontweight='bold')
ax_bar.text(1.5, n_incorrect / n_trials + (n_correct / n_trials) / 2, 'correct', rotation=90, va='center', ha='left', fontsize=9, color='blue', fontweight='bold')

ax_raster.axvline(0, color='black', linestyle='--', linewidth=0.8)
ax_raster.set_ylabel('Sorted Trial Number')
ax_raster.set_xlabel('T from Stim On (s)')

plt.tight_layout()
plt.show()

/opt/hostedtoolcache/Python/3.13.15/x64/lib/python3.13/site-packages/pynapple/core/time_index.py:151: UserWarning: timestamps are not sorted
  warn("timestamps are not sorted", UserWarning)
/opt/hostedtoolcache/Python/3.13.15/x64/lib/python3.13/site-packages/pynapple/core/base_class.py:56: UserWarning: 1 row(s) with NaN start/end time(s) were dropped.
  self.time_support = IntervalSet(start=self.index[0], end=self.index[-1])
/opt/hostedtoolcache/Python/3.13.15/x64/lib/python3.13/site-packages/pynapple/core/base_class.py:58: RuntimeWarning: divide by zero encountered in scalar divide
  self.rate = self.index.shape[0] / np.sum(

# Trial-aligned wheel velocity sorted by choice (left/right), aligned to first movement onset

# Masks for left/right choices
left_choice_mask = trials_df['mouse_wheel_choice'].values == 'counter_clockwise'
right_choice_mask = trials_df['mouse_wheel_choice'].values == 'clockwise'

# Get wheel velocity (in rad/s)
wheel_velocity = data['WheelVelocitySmoothed']

# Interpolate to regular grid (required by pynapple's compute_perievent)
rate = 1.0 / np.median(np.diff(wheel_velocity.t))
n_points = int(((wheel_velocity.t[-1] - wheel_velocity.t[0]) * rate)) + 1
regular_t = nap.Ts(t=np.linspace(wheel_velocity.t[0], wheel_velocity.t[-1], n_points))
wheel_velocity = wheel_velocity.interpolate(regular_t)

# Create Ts objects for first movement times, separated by choice (filter NaN)
# Restrict to wheel_velocity's time_support so perievent alignment works
ep = wheel_velocity.time_support
all_move = trials_df['wheel_movement_onset_time'].dropna().values
left_move = trials_df.loc[left_choice_mask, 'wheel_movement_onset_time'].dropna().values
right_move = trials_df.loc[right_choice_mask, 'wheel_movement_onset_time'].dropna().values
movement_times = nap.Ts(t=all_move, time_support=ep)
left_movement_times = nap.Ts(t=left_move, time_support=ep)
right_movement_times = nap.Ts(t=right_move, time_support=ep)

minmax = (-0.5, 1.0)

perievent_all = nap.compute_perievent(wheel_velocity, movement_times, window=minmax)
perievent_left = nap.compute_perievent(wheel_velocity, left_movement_times, window=minmax)
perievent_right = nap.compute_perievent(wheel_velocity, right_movement_times, window=minmax)

# Compute mean and SEM
left_mean = np.nanmean(perievent_left.values, axis=1)
left_sem = np.nanstd(perievent_left.values, axis=1) / np.sqrt(np.sum(~np.isnan(perievent_left.values), axis=1))
right_mean = np.nanmean(perievent_right.values, axis=1)
right_sem = np.nanstd(perievent_right.values, axis=1) / np.sqrt(np.sum(~np.isnan(perievent_right.values), axis=1))

time = perievent_left.t

# Sort trials by choice for raster (left first, then right)
sorted_indices = np.argsort(right_choice_mask)  # False (left) first, then True (right)
perievent_sorted = perievent_all.values[:, sorted_indices]

# Create figure
fig, axes = plt.subplots(2, 1, figsize=(4, 8), gridspec_kw={'height_ratios': [1, 3]}, sharex=True)

ax_avg = axes[0]
ax_avg.plot(time, left_mean, color='green', linewidth=1.5, label='CCW (left)')
ax_avg.fill_between(time, left_mean - left_sem, left_mean + left_sem, color='green', alpha=0.3)
ax_avg.plot(time, right_mean, color='olive', linewidth=1.5, label='CW (right)')
ax_avg.fill_between(time, right_mean - right_sem, right_mean + right_sem, color='olive', alpha=0.3)
ax_avg.axvline(0, color='black', linestyle='--', linewidth=0.8)
ax_avg.axhline(0, color='gray', linestyle='-', linewidth=0.5)
ax_avg.set_ylabel('Velocity (rad/s)')
ax_avg.set_title('Wheel velocity', fontsize=12, fontweight='bold', loc='center')
ax_avg.legend(loc='upper left', fontsize=8, framealpha=0.9)

ax_raster = axes[1]
n_trials = perievent_sorted.shape[1]
n_left = left_choice_mask.sum()
n_right = right_choice_mask.sum()

im = ax_raster.imshow(perievent_sorted.T, aspect='auto', cmap='gray_r',
                       extent=[time[0], time[-1], 0, n_trials],
                       vmin=np.nanpercentile(perievent_sorted, 5),
                       vmax=np.nanpercentile(perievent_sorted, 95),
                       origin='lower')

# Side color bar for left/right choice
bar_width = 0.05
ax_bar = ax_raster.inset_axes([1.02, 0, bar_width, 1])
ax_bar.axhspan(0, n_left / n_trials, color='green', alpha=0.8)
ax_bar.axhspan(n_left / n_trials, 1, color='olive', alpha=0.8)
ax_bar.set_xlim(0, 1)
ax_bar.set_ylim(0, 1)
ax_bar.axis('off')
ax_bar.text(1.5, (n_left / n_trials) / 2, 'CCW', rotation=90, va='center', ha='left', fontsize=9, color='green', fontweight='bold')
ax_bar.text(1.5, n_left / n_trials + (n_right / n_trials) / 2, 'CW', rotation=90, va='center', ha='left', fontsize=9, color='olive', fontweight='bold')

ax_raster.axvline(0, color='black', linestyle='--', linewidth=0.8)
ax_raster.set_ylabel('Sorted Trial Number')
ax_raster.set_xlabel('T from First Move (s)')

plt.tight_layout()
plt.show()

# Trial-aligned lick counts with correct/incorrect trial separation, aligned to feedback time

# Get lick timestamps from processing module
lick_events = nwbfile_processed.processing['lick_times']['EventsLickTimes']
lick_times = nap.Ts(t=lick_events.timestamps[:])

# Create Ts objects for feedback times
feedback_times = nap.Ts(t=trials_df['feedback_time'].values)
correct_feedback_times = nap.Ts(t=trials_df.loc[correct_mask, 'feedback_time'].values)
incorrect_feedback_times = nap.Ts(t=trials_df.loc[~correct_mask, 'feedback_time'].values)

minmax = (-0.5, 1.0)
bin_size = 0.02  # 20ms bins for counting

# Use compute_perievent to align lick timestamps around feedback
perievent_licks = nap.compute_perievent(lick_times, feedback_times, window=minmax)

# For the average trace, we need to count licks in bins
perievent_correct_licks = nap.compute_perievent(lick_times, correct_feedback_times, window=minmax)
perievent_incorrect_licks = nap.compute_perievent(lick_times, incorrect_feedback_times, window=minmax)

# Count licks in bins for average traces
correct_counts = perievent_correct_licks.count(bin_size)
incorrect_counts = perievent_incorrect_licks.count(bin_size)

# Mean across trials (columns are trials in the TsGroup)
correct_mean = np.nanmean(correct_counts.values, axis=1)
correct_sem = np.nanstd(correct_counts.values, axis=1) / np.sqrt(correct_counts.shape[1])
incorrect_mean = np.nanmean(incorrect_counts.values, axis=1)
incorrect_sem = np.nanstd(incorrect_counts.values, axis=1) / np.sqrt(incorrect_counts.shape[1])

time_counts = correct_counts.t

# Sort trials by outcome
sorted_trial_indices = np.argsort(correct_mask)

# Create figure
fig, axes = plt.subplots(2, 1, figsize=(4, 8), gridspec_kw={'height_ratios': [1, 3]}, sharex=True)

ax_avg = axes[0]
ax_avg.plot(time_counts, correct_mean, color='blue', linewidth=1.5, label='Correct')
ax_avg.fill_between(time_counts, correct_mean - correct_sem, correct_mean + correct_sem, color='blue', alpha=0.3)
ax_avg.plot(time_counts, incorrect_mean, color='red', linewidth=1.5, label='Incorrect')
ax_avg.fill_between(time_counts, incorrect_mean - incorrect_sem, incorrect_mean + incorrect_sem, color='red', alpha=0.3)
ax_avg.axvline(0, color='black', linestyle='--', linewidth=0.8)
ax_avg.set_ylabel('Licks (count)')
ax_avg.set_title('Licks', fontsize=12, fontweight='bold', loc='center')
ax_avg.legend(loc='upper left', fontsize=8, framealpha=0.9)

# Raster plot - plot individual lick events
ax_raster = axes[1]

# Plot licks as scatter points for each trial
n_trials = len(trials_df)
n_incorrect = (~correct_mask).sum()
n_correct = correct_mask.sum()

# Get lick times for each trial and plot as raster
for i, sorted_index in enumerate(sorted_trial_indices):
    if sorted_index in perievent_licks.keys():
        trial_licks = perievent_licks[sorted_index]
        if len(trial_licks) > 0:
            ax_raster.scatter(trial_licks.t, np.full(len(trial_licks), i),
                            c='black', s=1, marker='|', linewidths=0.5)

# Side color bar
bar_width = 0.05
ax_bar = ax_raster.inset_axes([1.02, 0, bar_width, 1])
ax_bar.axhspan(0, n_incorrect / n_trials, color='red', alpha=0.8)
ax_bar.axhspan(n_incorrect / n_trials, 1, color='blue', alpha=0.8)
ax_bar.set_xlim(0, 1)
ax_bar.set_ylim(0, 1)
ax_bar.axis('off')
ax_bar.text(1.5, (n_incorrect / n_trials) / 2, 'incorrect', rotation=90, va='center', ha='left', fontsize=9, color='red', fontweight='bold')
ax_bar.text(1.5, n_incorrect / n_trials + (n_correct / n_trials) / 2, 'correct', rotation=90, va='center', ha='left', fontsize=9, color='blue', fontweight='bold')

ax_raster.axvline(0, color='black', linestyle='--', linewidth=0.8)
ax_raster.set_ylabel('Sorted Trial Number')
ax_raster.set_xlabel('T from Feedback (s)')
ax_raster.set_xlim(minmax[0], minmax[1])
ax_raster.set_ylim(0, n_trials)

plt.tight_layout()
plt.show()

Data type	int16
Shape	(164362247, 384)
Array size	117.56 GiB
Chunk shape	(13020, 384)
Compression	gzip
Compression opts	4
Uncompressed size (bytes)	126230205696
Compressed size (bytes)	61453524661
Compression ratio	2.0540759279851195

Data type	float64
Shape	(164362247,)
Array size	1.22 GiB
Chunk shape	(1250000,)
Compression	gzip
Compression opts	4
Uncompressed size (bytes)	1314897976
Compressed size (bytes)	793462428
Compression ratio	1.6571647624378756

	location	group	electrode_name	contact_shapes	channel_name	rel_x	rel_y	probe_name	contact_ids	adc_group	adc_sample_order	group_name	x	y	z	imp
id
0	Infralimbic area layer 1	Probe01 pynwb.ecephys.ElectrodeGroup at 0x140701196667792\nFields:\n description: Electrode group for probe01\n device: Probe01 pynwb.device.Device at 0x140701196667472\nFields:\n description: A PRB_1_4_0480_1 probe. Additional metadata: {"model_name": "PRB_1_4_0480_1", "manufacturer": "imec", "port": "4", "slot": "2", "part_number": "PRB_1_4_0480_1"}\n manufacturer: imec\n serial_number: 18194814141\n\n location: Unresolved\n	e0	square	AP0,LF0	16.0	0.0	Probe01	e0	0	0	Probe01	3401.0	4105.0	5751.0	NaN
1	Infralimbic area layer 1	Probe01 pynwb.ecephys.ElectrodeGroup at 0x140701196667792\nFields:\n description: Electrode group for probe01\n device: Probe01 pynwb.device.Device at 0x140701196667472\nFields:\n description: A PRB_1_4_0480_1 probe. Additional metadata: {"model_name": "PRB_1_4_0480_1", "manufacturer": "imec", "port": "4", "slot": "2", "part_number": "PRB_1_4_0480_1"}\n manufacturer: imec\n serial_number: 18194814141\n\n location: Unresolved\n	e1	square	AP1,LF1	48.0	0.0	Probe01	e1	1	0	Probe01	3401.0	4105.0	5751.0	NaN
2	Infralimbic area layer 1	Probe01 pynwb.ecephys.ElectrodeGroup at 0x140701196667792\nFields:\n description: Electrode group for probe01\n device: Probe01 pynwb.device.Device at 0x140701196667472\nFields:\n description: A PRB_1_4_0480_1 probe. Additional metadata: {"model_name": "PRB_1_4_0480_1", "manufacturer": "imec", "port": "4", "slot": "2", "part_number": "PRB_1_4_0480_1"}\n manufacturer: imec\n serial_number: 18194814141\n\n location: Unresolved\n	e2	square	AP2,LF2	0.0	20.0	Probe01	e2	0	1	Probe01	3401.0	4087.0	5745.0	NaN
3	Infralimbic area layer 1	Probe01 pynwb.ecephys.ElectrodeGroup at 0x140701196667792\nFields:\n description: Electrode group for probe01\n device: Probe01 pynwb.device.Device at 0x140701196667472\nFields:\n description: A PRB_1_4_0480_1 probe. Additional metadata: {"model_name": "PRB_1_4_0480_1", "manufacturer": "imec", "port": "4", "slot": "2", "part_number": "PRB_1_4_0480_1"}\n manufacturer: imec\n serial_number: 18194814141\n\n location: Unresolved\n	e3	square	AP3,LF3	32.0	20.0	Probe01	e3	1	1	Probe01	3401.0	4087.0	5745.0	NaN

Data type	int16
Shape	(13696854, 384)
Array size	9.80 GiB
Chunk shape	(13020, 384)
Compression	gzip
Compression opts	4
Uncompressed size (bytes)	10519183872
Compressed size (bytes)	5364846101
Compression ratio	1.9607615342477835

Data type	float64
Shape	(13696854,)
Array size	104.50 MiB
Chunk shape	(1250000,)
Compression	gzip
Compression opts	4
Uncompressed size (bytes)	109574832
Compressed size (bytes)	71276130
Compression ratio	1.5373285839172244

Installing requirements¶

Locate Examples¶

Raw NWB file¶

root (NWBFile)

Locating and Accessing the AP and LF bands¶

ElectricalSeriesProbe00AP (ElectricalSeries)

ElectricalSeriesProbe00LF (ElectricalSeries)

Electrode Metadata¶

electrodes (ElectrodesTable)

Anatomical Location¶

localization (Localization)

Anatomical Location of Probe Structure¶

Anatomical Location of Probes in the CCFv3 space¶

Accessing Video Data¶

Processed files¶

root (NWBFile)

subject (IblSubject)

Epochs (Task vs Passive)¶

epochs (TimeIntervals)

Loading Spike Sorting Data¶

units (Units)

Accessing Waveform Data¶

Loading Trials Data¶

trials (TimeIntervals)

Overview of Trials¶

Psychometric Curves¶

Reaction Time by Contrast¶

Loading Wheel Data¶

wheel (ProcessingModule)

WheelPosition (SpatialSeries)

Loading Pose Estimation Data¶

pose_estimation (ProcessingModule)

LeftCamera (PoseEstimation)

Overlay Pose on Video¶

Loading Passive Data¶

passive_protocol (ProcessingModule)

gabor_table (TimeIntervals)

passive_intervals (TimeIntervals)

passive_task_replay (TimeIntervals)

Trial-Aligned Behavioral Data¶

Data type	uint64
Shape	(2392,)
Array size	18.69 KiB
Chunk shape	None
Compression	None
Compression opts	None
Uncompressed size (bytes)	19136
Compressed size (bytes)	19136
Compression ratio	1.0

Data type	uint8
Shape	(0, 0, 0)
Array size	0.00 bytes
Chunk shape	None
Compression	None
Compression opts	None
Uncompressed size (bytes)	0
Compressed size (bytes)	0
Compression ratio	undefined

Data type	object
Shape	(1,)
Array size	8.00 bytes
Chunk shape	None
Compression	None
Compression opts	None
Uncompressed size (bytes)	8
Compressed size (bytes)	16
Compression ratio	0.5

	start_time	stop_time	protocol_type	epoch_description
id
0	0.00000	4173.279130	task	Active behavioral task period. The mouse performs a decision-making task where it must turn a wheel to move a visual stimulus (Gabor patch) to the center of the screen. Correct responses are rewarded with water; incorrect responses trigger white noise feedback. The trials table contains detailed timing and outcome data for each trial during this epoch.
1	4173.27913	5478.753002	passive	Passive stimulus replay period. Visual and auditory stimuli are presented while the mouse is head-fixed but not performing any task. This epoch includes: (1) replay of task-relevant stimuli (Gabor patches at various contrasts and positions), (2) sparse noise stimuli for receptive field mapping, and (3) natural movie clips. Used for characterizing sensory responses independent of task engagement.

	x	y	z	localized_entity	probe_name	atlas_id	beryl_location	cosmos_location	brain_region
id
0	12.0	1999.0	-3773.0	0	Probe01	707	ILA	Isocortex	ILA1
1	12.0	1999.0	-3773.0	1	Probe01	707	ILA	Isocortex	ILA1
2	6.0	1999.0	-3755.0	2	Probe01	707	ILA	Isocortex	ILA1
3	6.0	1999.0	-3755.0	3	Probe01	707	ILA	Isocortex	ILA1

	location	group	electrode_name	contact_shapes	channel_name	rel_x	rel_y	probe_name	contact_ids	shank_ids	adc_group	adc_sample_order	group_name	x	y	z	imp	filtering
id
384	Cortical amygdalar area posterior part lateral...	Probe00 pynwb.ecephys.ElectrodeGroup at 0x1407...	e0	square	AP0,LF0	16.0	0.0	Probe00	e0		0	0	Probe00	7800.0	7016.0	2539.0	NaN
385	Cortical amygdalar area posterior part lateral...	Probe00 pynwb.ecephys.ElectrodeGroup at 0x1407...	e1	square	AP1,LF1	48.0	0.0	Probe00	e1		1	0	Probe00	7800.0	7016.0	2539.0	NaN
386	Cortical amygdalar area posterior part lateral...	Probe00 pynwb.ecephys.ElectrodeGroup at 0x1407...	e2	square	AP2,LF2	0.0	20.0	Probe00	e2		0	1	Probe00	7800.0	6996.0	2544.0	NaN
387	Cortical amygdalar area posterior part lateral...	Probe00 pynwb.ecephys.ElectrodeGroup at 0x1407...	e3	square	AP3,LF3	32.0	20.0	Probe00	e3		1	1	Probe00	7800.0	6996.0	2544.0	NaN
388	Cortical amygdalar area posterior part lateral...	Probe00 pynwb.ecephys.ElectrodeGroup at 0x1407...	e4	square	AP4,LF4	16.0	40.0	Probe00	e4		0	2	Probe00	7800.0	6976.0	2548.0	NaN
...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...	...
763	Field CA3	Probe00 pynwb.ecephys.ElectrodeGroup at 0x1407...	e379	square	AP379,LF379	32.0	3780.0	Probe00	e379		31	9	Probe00	7925.0	3198.0	3002.0	NaN
764	Field CA3	Probe00 pynwb.ecephys.ElectrodeGroup at 0x1407...	e380	square	AP380,LF380	16.0	3800.0	Probe00	e380		30	10	Probe00	7925.0	3180.0	3011.0	NaN
765	Field CA3	Probe00 pynwb.ecephys.ElectrodeGroup at 0x1407...	e381	square	AP381,LF381	48.0	3800.0	Probe00	e381		31	10	Probe00	7925.0	3180.0	3011.0	NaN
766	Field CA3	Probe00 pynwb.ecephys.ElectrodeGroup at 0x1407...	e382	square	AP382,LF382	0.0	3820.0	Probe00	e382		30	11	Probe00	7925.0	3162.0	3020.0	NaN
767	Field CA3	Probe00 pynwb.ecephys.ElectrodeGroup at 0x1407...	e383	square	AP383,LF383	32.0	3820.0	Probe00	e383		31	11	Probe00	7925.0	3162.0	3020.0	NaN

	probe_name	location	x	y	z	rel_x	rel_y
id
257	Probe01	Secondary motor area layer 2/3	3075.0	1770.0	5100.0	48.0	2560.0
713	Probe00	fiber tracts	7925.0	3729.0	2901.0	48.0	3280.0
234	Probe01	Secondary motor area layer 5	3075.0	1989.0	5142.0	0.0	2340.0
403	Probe00	Cortical amygdalar area posterior part lateral...	7800.0	6838.0	2578.0	32.0	180.0
459	Probe00	Basomedial amygdalar nucleus posterior part	7800.0	6278.0	2656.0	32.0	740.0

	x	y	z	localized_entity	probe_name	brain_region
id
222	3084.0	2094.0	5184.0	location \ id ...	Probe01	MOs5
103	3251.0	3174.0	5495.0	location \ id ...	Probe01	PL2/3
62	3276.0	3545.0	5607.0	location \ id ...	Probe01	ILA1
328	3015.0	1071.0	4882.0	location ...	Probe01	void
393	7800.0	6937.0	2557.0	...	Probe00	COApl

	start_time	stop_time	position	contrast	phase
id
0	5126.874515	5127.185221	-35.0	0.0625	3.124994
1	5129.124614	5129.435687	-35.0	1.0000	2.162905
2	5131.441141	5131.768878	35.0	1.0000	2.200607
3	5134.424608	5134.735413	35.0	0.0625	1.741850

	start_time	stop_time	protocol_name
id
0	4185.545797	4785.626265	spontaneousActivity
1	5111.224793	5478.753002	taskReplay

	start_time	stop_time	stim_type
id
0	5125.087475	5125.587463	noise
1	5127.316809	5127.416800	tone
2	5130.037133	5130.137124	tone
3	5130.516457	5130.583451	valve

	start_time	stop_time	peak_amplitude
id
0	0.343596	0.723596	0.132027
1	3.532596	4.394596	0.700072
2	6.731596	11.317596	11.646018
3	22.972596	24.000596	0.591526

	start_time	stop_time	quiescence_period	gabor_stimulus_onset_time	auditory_cue_time	wheel_movement_onset_time	choice_registration_time	feedback_time	gabor_stimulus_offset_time	gabor_stimulus_contrast	gabor_stimulus_side	mouse_wheel_choice	is_mouse_rewarded	reward_volume_uL	probability_left	block_type	block_index
id
0	103.884959	115.226015	0.541949	NaN	113.439040	113.477967	113.677692	113.677785	114.725958	100.0	right	counter_clockwise	True	1.5	0.5	unbiased	0
1	115.866423	119.975890	0.457997	116.381644	116.382710	117.707967	118.402564	118.402663	119.475868	25.0	left	clockwise	True	1.5	0.5	unbiased	0
2	120.616099	123.259428	0.426862	121.098089	121.099023	121.328967	121.695111	121.695203	122.759341	25.0	right	counter_clockwise	True	1.5	0.5	unbiased	0
3	123.949635	131.142860	0.423191	124.431558	124.432692	128.465967	128.570015	128.571022	130.642804	12.5	right	clockwise	False	0.0	0.5	unbiased	0

	device	ml	ap	dv	depth_um	theta	phi	pid	trajectory_source	roll
id
0	Probe00 pynwb.device.Device at 0x140700232784496\nFields:\n description: A PRB_1_4_0480_1 probe. Additional metadata: {"model_name": "PRB_1_4_0480_1", "manufacturer": "imec", "port": "3", "slot": "2", "part_number": "PRB_1_4_0480_1"}\n manufacturer: imec\n serial_number: 18005102151\n	-2464.0	-2600.0	-143.0	6585.415277	6.353030	-18.593623	1ca6cd06-1ed5-45af-b73a-017d5e7cff48	Ephys aligned histology track	0.0
1	Probe00 pynwb.device.Device at 0x140700232784496\nFields:\n description: A PRB_1_4_0480_1 probe. Additional metadata: {"model_name": "PRB_1_4_0480_1", "manufacturer": "imec", "port": "3", "slot": "2", "part_number": "PRB_1_4_0480_1"}\n manufacturer: imec\n serial_number: 18005102151\n	-1900.0	-2150.0	0.0	6400.000000	20.000000	0.000000	1ca6cd06-1ed5-45af-b73a-017d5e7cff48	Planned	0.0
2	Probe00 pynwb.device.Device at 0x140700232784496\nFields:\n description: A PRB_1_4_0480_1 probe. Additional metadata: {"model_name": "PRB_1_4_0480_1", "manufacturer": "imec", "port": "3", "slot": "2", "part_number": "PRB_1_4_0480_1"}\n manufacturer: imec\n serial_number: 18005102151\n	-2090.0	-2100.0	0.0	6300.000000	20.000000	0.000000	1ca6cd06-1ed5-45af-b73a-017d5e7cff48	Micro-manipulator	0.0
3	Probe00 pynwb.device.Device at 0x140700232784496\nFields:\n description: A PRB_1_4_0480_1 probe. Additional metadata: {"model_name": "PRB_1_4_0480_1", "manufacturer": "imec", "port": "3", "slot": "2", "part_number": "PRB_1_4_0480_1"}\n manufacturer: imec\n serial_number: 18005102151\n	-2289.0	-2550.0	-93.0	6888.922029	8.106919	-9.751776	1ca6cd06-1ed5-45af-b73a-017d5e7cff48	Histology track	0.0

	device	ml	ap	dv	depth_um	theta	phi	pid	trajectory_source	roll
id
0	Probe01 pynwb.device.Device at 0x140700232784800\nFields:\n description: A PRB_1_4_0480_1 probe. Additional metadata: {"model_name": "PRB_1_4_0480_1", "manufacturer": "imec", "port": "4", "slot": "2", "part_number": "PRB_1_4_0480_1"}\n manufacturer: imec\n serial_number: 18194814141\n	-739.0	2350.0	-1143.0	2757.103993	17.385231	158.599449	19c5b0d5-a255-47ff-9f8d-639e634a7b61	Ephys aligned histology track	0.0
1	Probe01 pynwb.device.Device at 0x140700232784800\nFields:\n description: A PRB_1_4_0480_1 probe. Additional metadata: {"model_name": "PRB_1_4_0480_1", "manufacturer": "imec", "port": "4", "slot": "2", "part_number": "PRB_1_4_0480_1"}\n manufacturer: imec\n serial_number: 18194814141\n	-700.0	2200.0	0.0	2500.000000	19.000000	180.000000	19c5b0d5-a255-47ff-9f8d-639e634a7b61	Planned	0.0
2	Probe01 pynwb.device.Device at 0x140700232784800\nFields:\n description: A PRB_1_4_0480_1 probe. Additional metadata: {"model_name": "PRB_1_4_0480_1", "manufacturer": "imec", "port": "4", "slot": "2", "part_number": "PRB_1_4_0480_1"}\n manufacturer: imec\n serial_number: 18194814141\n	-650.0	2200.0	0.0	3000.000000	19.000000	180.000000	19c5b0d5-a255-47ff-9f8d-639e634a7b61	Micro-manipulator	0.0
3	Probe01 pynwb.device.Device at 0x140700232784800\nFields:\n description: A PRB_1_4_0480_1 probe. Additional metadata: {"model_name": "PRB_1_4_0480_1", "manufacturer": "imec", "port": "4", "slot": "2", "part_number": "PRB_1_4_0480_1"}\n manufacturer: imec\n serial_number: 18194814141\n	-739.0	2375.0	-1168.0	3198.109372	17.581281	154.880882	19c5b0d5-a255-47ff-9f8d-639e634a7b61	Histology track	0.0

	unit_name	cluster_uuid	spike_count	firing_rate	kilosort2_label	ibl_quality_score
id
678	probe00_678	5c4ebfea-548d-4985-9d67-0300ba002de2	749	0.136710	good	0.333333
84	probe00_84	5db3fbd7-563c-4356-abed-a4914bac42d4	18373	3.353506	mua	0.000000
386	probe00_386	05f2c1e2-e99d-41bf-ac62-91b483964a0e	78710	14.366430	good	0.666667
137	probe00_137	06b20157-ccde-4d76-a49b-72746303ae7f	1077	0.196578	good	0.666667
82	probe00_82	7b748b2f-0256-40c8-805f-23d2928f0a02	40009	7.302586	mua	0.000000

	unit_name	median_spike_amplitude_uV	distance_from_probe_tip_um
id
513	probe00_513	57.913134	3240.0
308	probe00_308	36.973923	2100.0
518	probe00_518	107.732564	3280.0
195	probe00_195	53.754077	1020.0
630	probe00_630	31.791270	900.0