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

import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
from pynwb import NWBHDF5IO
from ndx_structured_behavior.plot import (
    plot_events,
    plot_actions,
    plot_states,
    compute_state_transition_matrix,
)
from dandi.dandiapi import DandiAPIClient
import remfile
import h5py

DANDISET_ID = "001550"
ASSET_PATH  = "sub-P100/sub-P100_ses-TaskSwitch6-190423a.nwb"

with DandiAPIClient() as client:
    asset = client.get_dandiset(DANDISET_ID, "draft").get_asset_by_path(ASSET_PATH)
    s3_url = asset.get_content_url(follow_redirects=1, strip_query=False)

file_system = remfile.File(s3_url)
h5_file = h5py.File(file_system, "r")
io = NWBHDF5IO(file=h5_file)
nwbfile = io.read()

print(f"Session ID         : {nwbfile.session_id}")
print(f"Session start time : {nwbfile.session_start_time}")
print(f"Subject ID         : {nwbfile.subject.subject_id}")
print(f"Experimenter       : {nwbfile.experimenter}")
print(f"Session description : {nwbfile.session_description}")

Session ID         : TaskSwitch6-190423a
Session start time : 2019-04-23 11:30:00+01:00
Subject ID         : P100
Experimenter       : ('Pagan, Marino', 'Tang, Vincent D.')
Session description : This session contains behavioral data from a rat performing a context-dependent decision-making task.
The task required the animal to accumulate sensory evidence over time and select an action based on the context
cue for that trial. This dataset includes event timings, trial structure, stimulus parameters, and subject responses.

print("Subject ID:    ", nwbfile.subject.subject_id)
print("Species:       ", nwbfile.subject.species)
print("Strain:        ", nwbfile.subject.strain)
print("Sex:           ", nwbfile.subject.sex)

Subject ID:     P100
Species:        Rattus norvegicus
Strain:         Long Evans
Sex:            M

task = nwbfile.lab_meta_data["task"]

print("=== Event types (port pokes) ===")
display(task.event_types[:])

print("=== Action types (sounds / outputs) ===")
display(task.action_types[:])

print("=== State types ===")
display(task.state_types[:])

=== Event types (port pokes) ===

=== Action types (sounds / outputs) ===

=== State types ===

task_recording = nwbfile.acquisition["task_recording"]
events       = task_recording.events
states       = task_recording.states
actions      = task_recording.actions
event_types  = task.event_types
state_types  = task.state_types
action_types = task.action_types

print("\n--- Events (first 5) ---")
display(
    pd.merge(
        events[:5], event_types[:],
        left_on="event_type", right_on="id",
    )
)

print("--- States (first 5) ---")
display(
    pd.merge(
        states[:5], state_types[:],
        left_on="state_type", right_on="id",
    )
)

print("--- Actions (first 5) ---")
display(
    pd.merge(
        actions[:5], action_types[:],
        left_on="action_type", right_on="id",
    )
)

--- Events (first 5) ---

--- States (first 5) ---

--- Actions (first 5) ---

fig = plot_events(
    events=events[20:100],
    event_types=event_types,
    show_event_values=True,
    figsize=(18, 4),
    marker_size=500,
)
plt.title(f"Events — {nwbfile.session_id}", fontsize=16)
plt.tight_layout()
plt.show()

fig = plot_actions(
    actions=actions[20:100],
    action_types=action_types,
    figsize=(18, 4),
    marker_size=500,
)
plt.title(f"Actions — {nwbfile.session_id}", fontsize=16)
plt.tight_layout()
plt.show()

plot_states(
    figsize=(14, 7),
    states=states[20:100],
    state_types=state_types,
    marker_size=500,
)
plt.title(f"States — {nwbfile.session_id}", fontsize=16)
plt.tight_layout()
plt.show()

trials = nwbfile.trials
print(f"Number of trials : {len(trials)}")
print(f"Columns          : {list(trials.colnames)}")

Number of trials : 900
Columns          : ['start_time', 'stop_time', 'states', 'events', 'actions', 'cpoke_start_time', 'left_hi', 'right_hi', 'left_lo', 'right_lo', 'crosstalk_dir', 'crosstalk_freq', 'bup_width', 'bup_ramp', 'vol_low', 'vol_hi', 'vol', 'gamma_dir', 'gamma_freq', 'duration', 'freq_lo', 'freq_hi', 'HistorySection_hit_history', 'HistorySection_side_history', 'HistorySection_quadrant_history', 'HistorySection_task_history', 'HistorySection_incoh_history', 'HistorySection_gammadir_history', 'HistorySection_gammafreq_history', 'HistorySection_result_history']

pd.set_option("display.max_columns", None)

# Exclude ragged / region-reference columns that require special handling
REGION_COLS = {"left_hi", "right_hi", "left_lo", "right_lo", "states", "events", "actions"}
scalar_cols = [c for c in trials.colnames if c not in REGION_COLS]

trials[:5][scalar_cols]

trials_df = trials[10:16]

fig = plt.figure(figsize=(18, 10))

# Events (blue ticks)
events_index = [j for i in trials_df["events"] for j in i]
if events_index:
    plot_events(events=events[events_index], event_types=event_types,
                show_event_values=True, marker_size=500, marker_width=2,
                marker_color="tab:blue", y_offset=0, fontsize=14, fig=fig)

# Actions (green ticks)
actions_index = [j for i in trials_df["actions"] for j in i]
y_offset = int(np.ceil(plt.ylim()[1]))
if y_offset == plt.ylim()[1]:
    y_offset += 1
if actions_index:
    plot_actions(actions=actions[actions_index], action_types=action_types,
                 show_action_values=True, marker_size=500, marker_width=2,
                 marker_color="tab:green", y_offset=y_offset, keep_yticks=True,
                 fontsize=14, fig=fig)

# States (steelblue rectangles — replaces hardcoded black)
states_index = [j for i in trials_df["states"] for j in i]
y_offset = int(np.ceil(plt.ylim()[1]))
if y_offset == plt.ylim()[1]:
    y_offset += 1
if states_index:
    plot_states(states=states[states_index], state_types=state_types,
                rectangle_color="steelblue", rectangle_height=1,
                marker_color="red", marker_size=500,
                y_offset=y_offset, keep_yticks=True, fontsize=14, fig=fig)

plt.title(f"Trial structure — {nwbfile.session_id} (trials 10–15)", fontsize=16)
plt.tight_layout()
plt.show()

STIMULUS_SCALAR_COLS = [
    "gamma_dir", "gamma_freq", "duration", "freq_lo", "freq_hi",
    "vol", "vol_low", "vol_hi", "bup_width", "bup_ramp",
    "crosstalk_dir", "crosstalk_freq",
]
PULSE_COLS = ["left_hi", "right_hi", "left_lo", "right_lo"]

has_stimulus = any(c in trials.colnames for c in STIMULUS_SCALAR_COLS)
has_pulses   = any(c in trials.colnames for c in PULSE_COLS)

print(f"Has stimulus cols : {has_stimulus}")
print(f"Has pulse cols    : {has_pulses}")

Has stimulus cols : True
Has pulse cols    : True

if has_stimulus:
    avail_scalar = [c for c in STIMULUS_SCALAR_COLS if c in trials.colnames]
    print(f"Stimulus scalar columns: {avail_scalar}")
    display(trials[:5][avail_scalar])
else:
    print("No stimulus columns in this session — skipping.")

Stimulus scalar columns: ['gamma_dir', 'gamma_freq', 'duration', 'freq_lo', 'freq_hi', 'vol', 'vol_low', 'vol_hi', 'bup_width', 'bup_ramp', 'crosstalk_dir', 'crosstalk_freq']

if has_pulses:
    trial_idx = 0
    cpoke_t = trials["cpoke_start_time"][trial_idx]
    print(f"Trial {trial_idx}  |  cpoke_start_time = {cpoke_t:.4f} s from session start")
    print()
    print(f"  {'column':12s}  {'n pulses':>8s}  {'first 4 (rel. to cpoke, s)':30s}  {'first 4 (abs. session time, s)'}")
    print(f"  {'-'*12}  {'-'*8}  {'-'*30}  {'-'*30}")
    for col in PULSE_COLS:
        if col not in trials.colnames:
            continue
        rel = np.array(trials[col][trial_idx])
        abs_ = rel + cpoke_t if not np.isnan(cpoke_t) else rel
        rel_preview  = str(np.round(rel[:4],  4).tolist()) + (" ..." if len(rel) > 4 else "")
        abs_preview  = str(np.round(abs_[:4], 4).tolist()) + (" ..." if len(abs_) > 4 else "")
        print(f"  {col:12s}  {len(rel):>8d}  {rel_preview:30s}  {abs_preview}")
else:
    print("No pulse-time columns in this session — skipping.")

Trial 0  |  cpoke_start_time = 56.1463 s from session start

  column        n pulses  first 4 (rel. to cpoke, s)      first 4 (abs. session time, s)
  ------------  --------  ------------------------------  ------------------------------
  left_hi             34  [0.0236, 0.1519, 0.1672, 0.1692] ...  [56.1699, 56.2981, 56.3135, 56.3155] ...
  right_hi            17  [0.0536, 0.0714, 0.3894, 0.3982] ...  [56.1999, 56.2177, 56.5357, 56.5444] ...
  left_lo              3  [0.039, 0.0563, 0.3522]         [56.1853, 56.2025, 56.4985]
  right_lo             4  [0.0184, 0.2052, 0.2583, 1.1511]  [56.1647, 56.3515, 56.4046, 57.2973]

if has_pulses:
    N_TRIALS = 10
    colors = {
        "left_hi":  "tab:blue",
        "right_hi": "tab:orange",
        "left_lo":  "tab:cyan",
        "right_lo": "tab:red",
    }

    fig, axes = plt.subplots(N_TRIALS, 1, figsize=(14, N_TRIALS * 0.9), sharex=True)

    for trial_idx, ax in enumerate(axes):
        # Pulse times are already relative to cpoke onset — plot directly so
        # t=0 is cpoke entry on every trial (natural alignment point).
        for col, color in colors.items():
            if col not in trials.colnames:
                continue
            pulses = np.array(trials[col][trial_idx])
            if len(pulses) == 0:
                continue
            ax.vlines(pulses, 0, 1, color=color, linewidth=1.2,
                      label=col if trial_idx == 0 else None)
        ax.axvline(0, color="gray", linewidth=0.8, linestyle="--")  # cpoke onset
        ax.set_ylabel(f"Trial#{trial_idx + 1}", fontsize=10, rotation=0, labelpad=25)
        ax.set_yticks([])
        ax.spines[["top", "right", "left"]].set_visible(False)

    axes[0].legend(loc="upper right", fontsize=9, ncol=4)
    axes[-1].set_xlabel("Time relative to cpoke onset (s)", fontsize=12)
    fig.suptitle(
        f"Auditory pulse raster — {nwbfile.session_id} (first {N_TRIALS} trials)\n"
        "t=0 = centre-port entry (cpoke onset)",
        fontsize=13, y=1.01,
    )
    plt.tight_layout()
    plt.show()
else:
    print("No pulse-time columns in this session — skipping raster plot.")

pd.set_option("display.max_rows", None)
pd.set_option("display.max_colwidth", 80)

task_args = task.task_arguments[:]
print(f"{len(task_args)} task arguments found")
task_args[["argument_name", "expression_type", "expression"]].head(20)

224 task arguments found

count_df, prob_df = compute_state_transition_matrix(states=states, state_types=state_types)

print("State transition probability matrix (top-left corner):")
display(prob_df.iloc[:10, :10])

State transition probability matrix (top-left corner):

import networkx as nx
from matplotlib.patches import Patch

# ── Build base active matrix (drop all-zero rows/cols) ────────────────────────
mask   = (prob_df != 0).any(axis=1)
active = prob_df.loc[mask, prob_df.columns[(prob_df != 0).any(axis=0)]]

# ── Filter out continuation states that end with "2" ─────────────────────────
drop_states = [s for s in active.index if s.endswith("2")]
filtered = active.drop(
    index=drop_states,
    columns=[c for c in drop_states if c in active.columns],
    errors="ignore",
)
m2 = (filtered != 0).any(axis=1)
filtered = filtered.loc[m2, filtered.columns[(filtered != 0).any(axis=0)]]

print(f"Dropped states : {drop_states}")
print(f"Remaining      : {filtered.index.tolist()}")

# ── Build directed graph ──────────────────────────────────────────────────────
G = nx.DiGraph()
for src in filtered.index:
    for dst in filtered.columns:
        prob = filtered.loc[src, dst]
        if prob > 0:
            G.add_edge(src, dst, weight=prob)

# ── Node colours by semantic role ─────────────────────────────────────────────
REWARD_KW = {"hit", "reward", "drink", "soft_drink"}
ERROR_KW  = {"error", "violation", "punish", "timeout", "pun", "warning", "danger"}

def node_color(name):
    if name == "state_0":                          return "#2ecc71"  # green
    if any(k in name for k in REWARD_KW):          return "#f1c40f"  # gold
    if any(k in name for k in ERROR_KW):           return "#e74c3c"  # red
    return "#aec6cf"                                                  # blue-grey

node_colors = [node_color(n) for n in G.nodes()]
node_sizes  = [3500 if n == "state_0" else 2200 for n in G.nodes()]

# ── Layout ────────────────────────────────────────────────────────────────────
try:
    pos = nx.kamada_kawai_layout(G, weight=None)
except Exception:
    pos = nx.spring_layout(G, seed=42, k=2.0)

# ── Draw ──────────────────────────────────────────────────────────────────────
fig, ax = plt.subplots(figsize=(16, 11))

# 1. Nodes first
nx.draw_networkx_nodes(
    G, pos, ax=ax,
    node_color=node_colors, node_size=node_sizes,
    edgecolors="white", linewidths=1.5, alpha=0.95,
)

# 2. Edges — width and alpha scale with probability
all_weights = [d["weight"] for _, _, d in G.edges(data=True)]
max_w = max(all_weights) if all_weights else 1.0

for src, dst, data in G.edges(data=True):
    w = data["weight"]
    nx.draw_networkx_edges(
        G, pos, edgelist=[(src, dst)], ax=ax,
        width=0.6 + 4.4 * w / max_w,
        alpha=0.25 + 0.70 * w / max_w,
        edge_color="steelblue",
        arrows=True,
        arrowstyle="-|>",
        arrowsize=14,
        connectionstyle="arc3,rad=0.12",
        min_source_margin=22,
        min_target_margin=22,
    )

# 3. Edge labels (dominant transitions only)
LABEL_THRESHOLD = 0.25
edge_labels = {
    (src, dst): f"{data['weight']:.2f}"
    for src, dst, data in G.edges(data=True)
    if data["weight"] >= LABEL_THRESHOLD
}
nx.draw_networkx_edge_labels(
    G, pos, edge_labels=edge_labels, ax=ax,
    font_size=9, font_color="#222222",
    bbox=dict(boxstyle="round,pad=0.15", fc="white", alpha=0.65),
)

# 4. Node labels LAST — drawn on top of everything with an opaque white backing
#    so edges passing through a node center cannot bleed into the text.
for node, (x, y) in pos.items():
    ax.text(
        x, y, node,
        ha="center", va="center",
        fontsize=10, fontweight="bold", color="black",
        bbox=dict(boxstyle="round,pad=0.25", fc="white", ec="none", alpha=0.75),
        zorder=5,
    )

# ── Legend ────────────────────────────────────────────────────────────────────
legend_elements = [
    Patch(facecolor="#2ecc71", label="state_0  (trial start)"),
    Patch(facecolor="#f1c40f", label="Reward / success"),
    Patch(facecolor="#e74c3c", label="Error / punishment"),
    Patch(facecolor="#aec6cf", label="Normal flow"),
]
ax.legend(handles=legend_elements, loc="lower left", fontsize=10, framealpha=0.85)
ax.text(
    0.99, 0.01,
    f"Edge labels: p ≥ {LABEL_THRESHOLD}  |  States ending in '2' removed",
    transform=ax.transAxes, ha="right", va="bottom", fontsize=8, color="gray",
)
ax.set_title(f"State transition graph — {nwbfile.session_id}", fontsize=15, pad=12)
ax.axis("off")
plt.tight_layout()
plt.show()

Dropped states : ['wait_for_cpoke_wait2', 'hit_state2', 'nic_error_state2']
Remaining      : ['state_0', 'sending_trialnum', 'check_next_trial_ready', 'wait_for_cpoke_dir', 'wait_for_cpoke_freq', 'wait_for_cpoke_bis', 'nic_prestim', 'cpoke', 'cpoke_in', 'cpoke_out', 'wait_for_cout', 'wait_for_spoke', 'error_state', 'timeout_state', 'clean_up_state']

	action_name
id
0	direct_reward
1	stimulatorwave
2	cpoke_timer

	timestamp	event_type	value	duration	event_name
0	52.709765	2	out	0.024506	R
1	52.757265	2	out	0.048500	R
2	52.817280	2	out	0.058997	R
3	52.890275	2	out	0.054990	R
4	52.951767	2	out	0.067498	R

	start_time	stop_time	state_type	state_name
0	52.710269	59.287265	0	state_0
1	52.710269	52.726272	1	sending_trialnum
2	52.726272	53.726775	3	wait_for_cpoke
3	53.726775	53.727276	4	wait_for_cpoke_wait2
4	53.727276	54.727770	5	wait_for_cpoke_dir

	timestamp	action_type	value	duration	action_name
0	56.146765	1	out	0.010000	stimulatorwave
1	56.157265	1	out	0.009501	stimulatorwave
2	57.546765	2	out	0.000000	cpoke_timer
3	64.004271	1	out	0.010004	stimulatorwave
4	64.014766	1	out	0.009999	stimulatorwave

	start_time	stop_time	cpoke_start_time	crosstalk_dir	crosstalk_freq	bup_width	bup_ramp	vol_low	vol_hi	vol	gamma_dir	gamma_freq	duration	freq_lo	freq_hi	HistorySection_hit_history	HistorySection_side_history	HistorySection_quadrant_history	HistorySection_task_history	HistorySection_incoh_history	HistorySection_gammadir_history	HistorySection_gammafreq_history	HistorySection_result_history
id
0	52.710269	59.287265	56.146274	0	0	5	2	1	1	0.15	-1.0	2.5	1.3	6500	14200	NaN	l	3	d	1	-1.0	2.5	3
1	59.287771	67.177265	64.003770	0	0	5	2	1	1	0.15	-4.0	1.0	1.3	6500	14200	NaN	l	3	d	1	-4.0	1.0	3
2	67.177770	118.347265	115.138290	0	0	5	2	1	1	0.15	4.0	-1.0	1.3	6500	14200	NaN	r	1	d	1	4.0	-1.0	3
3	118.347765	123.885765	120.366770	0	0	5	2	1	1	0.15	2.5	-4.0	1.3	6500	14200	NaN	r	1	d	1	2.5	-4.0	3
4	123.886280	129.112281	125.904278	0	0	5	2	1	1	0.15	1.0	1.0	1.3	6500	14200	NaN	r	4	d	0	1.0	1.0	3

Installing requirements¶

Pagan Lab — Behavioral Data¶

1. Stream the NWB file from DANDI¶

2. Session and subject metadata¶

3. Task metadata¶

4. Behavioral data — Events, States, Actions¶

Events¶

States¶

Actions¶

Plot events, states, and actions¶

5. Trials table¶

Plot trial structure (states + events + actions per trial)¶

6. Protocol-specific: Stimulus data (TaskSwitch family only)¶

Pulse-time reference frame¶

Pulse-time raster (first 10 trials)¶

7. Task arguments¶

8. State transition analysis¶

	state_name
id
0	state_0
1	sending_trialnum
2	check_next_trial_ready
3	wait_for_cpoke
4	wait_for_cpoke_wait2
5	wait_for_cpoke_dir
6	wait_for_cpoke_freq
7	wait_for_cpoke_bis
8	nic_prestim
9	cpoke
10	cpoke_in
11	cpoke_out
12	wait_for_cout
13	wait_for_spoke
14	hit_state
15	hit_state2
16	error_state
17	wait_state
18	nic_error_state
19	nic_error_state2
20	timeout_state
21	clean_up_state

	argument_name	expression_type	expression
id
0	ProtocolsSection_n_done_trials	integer	900
1	ProtocolsSection_n_started_trials	integer	901
2	ProtocolsSection_n_completed_trials	integer	900
3	TaskSwitch6_sessid	integer	689581
4	SavingSection_data_file	string	C:\ratter\SoloData\Data\Marino\P100\data_@TaskSwitch6_Marino_P100_190423_ASV...
5	SavingSection_settings_file_load_time	integer	0
6	SavingSection_experimenter	string	Marino
7	SavingSection_ratname	string	P100
8	SavingSection_hostname	string	localhost
9	SavingSection_SaveTime	string	23-Apr-2019 13:52:45
10	SavingSection_n_autosave_calls	integer	902
11	SavingSection_autosave_frequency	integer	20
12	SavingSection_save_all_data_to_sql	integer	0
13	SavingSection_title	string	SavingSection
14	SavingSection_interactive_by_default	integer	1
15	CommentsSection_CommentsShow	integer	0
16	WaterValvesSection_WaterShow	integer	0
17	WaterValvesSection_Left_volume	integer	24
18	WaterValvesSection_LeftWValveTime	float	0.051555780933062886
19	WaterValvesSection_Center_volume	integer	24

to	state_0	sending_trialnum	check_next_trial_ready	wait_for_cpoke	wait_for_cpoke_wait2	wait_for_cpoke_dir	wait_for_cpoke_freq	wait_for_cpoke_bis	nic_prestim	cpoke
from
state_0	0.000000	0.526667	0.0	0.473333	0.0	0.000000	0.000000	0.0	0.000000	0.000000
sending_trialnum	0.473333	0.000000	0.0	0.526667	0.0	0.000000	0.000000	0.0	0.000000	0.000000
check_next_trial_ready	0.487097	0.512903	0.0	0.000000	0.0	0.000000	0.000000	0.0	0.000000	0.000000
wait_for_cpoke	0.000000	0.000000	0.0	0.000000	1.0	0.000000	0.000000	0.0	0.000000	0.000000
wait_for_cpoke_wait2	0.000000	0.000000	0.0	0.000000	0.0	0.546563	0.453437	0.0	0.000000	0.000000
wait_for_cpoke_dir	0.000000	0.000000	0.0	0.000000	0.0	0.000000	0.000000	1.0	0.000000	0.000000
wait_for_cpoke_freq	0.000000	0.000000	0.0	0.000000	0.0	0.000000	0.000000	1.0	0.000000	0.000000
wait_for_cpoke_bis	0.000000	0.000000	0.0	0.000000	0.0	0.000000	0.000000	0.0	0.984479	0.000000
nic_prestim	0.000000	0.000000	0.0	0.002252	0.0	0.000000	0.000000	0.0	0.000000	0.997748
cpoke	0.000000	0.000000	0.0	0.000000	0.0	0.000000	0.000000	0.0	0.000000	0.000000