"""Set of functions for plotting paw state data.
Adapted from https://github.com/rgs2151/plumber/blob/main/plumber/pipes/ibl_overview.py.
"""
import logging
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import numpy as np
import pandas as pd
import seaborn as sns
from matplotlib.colors import ListedColormap
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
from scipy.interpolate import interp1d
from brainbox.behavior.wheel import interpolate_position, velocity_filtered
logger = logging.getLogger('ibllib')
# Global constants
MODEL_TO_USE = 'e_mode'
SILENCE_WINDOW = 5
FRAME_COUNTER_DURATION = 10
L_THRESH = 0.0
COLORS = ['#DD8452', '#4C72B0', '#55A868', '#CC78BC']
ORIG_LABELS = ('background', 'still', 'move', 'wheel_turn', 'groom')
STATE_LABELS = ['Still', 'Move', 'Wheel Turn', 'Groom']
STATE_KEYS = {1: 'Still', 2: 'Move', 3: 'Wheel Turn', 4: 'Groom'}
STATE_MAP = np.vectorize(lambda x: STATE_KEYS[x])
# -------------------------------------------------
# Data extraction functions
# -------------------------------------------------
[docs]
def gen_data_df(marker_data, probs, ens_vars, cam_times):
"""
Generate processed data DataFrame with state predictions and ensemble statistics.
:param marker_data: DataFrame with marker positions and velocities
:param probs: DataFrame of state probabilities
:param ens_vars: DataFrame of ensemble variances
:param cam_times: Camera timestamps
:returns: Tuple of (processed_data_df, transition_frame_indices)
"""
# Calculate ensemble mode and variance
ensemble_mode = probs.to_numpy().argmax(axis=1)
idxs_good = probs.isna().sum(axis=1) == 0
ensemble_variance = ens_vars.sum(axis=1).to_numpy()
# Remove unfilled final batch
data_df = marker_data.copy().loc[idxs_good]
ensemble_mode = ensemble_mode[idxs_good]
ensemble_variance = ensemble_variance[idxs_good]
cam_times = cam_times[idxs_good]
# Calculate derived features
data_df['paw_speed'] = np.sqrt(data_df['paw_x_vel'] ** 2 + data_df['paw_y_vel'] ** 2)
data_df['wheel_speed'] = data_df['wheel_vel'].abs()
# Add predictions
data_df['e_mode'] = STATE_MAP(ensemble_mode)
data_df['e_var'] = ensemble_variance
data_df['frame_id'] = np.arange(0, len(data_df))
data_df['times'] = cam_times
# Find transition frames
any_transition_frames = data_df[
(data_df['e_mode'].shift() != data_df['e_mode'])
].index
# Silence frames around transitions
for noise_frame in any_transition_frames:
data_df.loc[noise_frame - SILENCE_WINDOW:noise_frame + SILENCE_WINDOW, 'e_var'] = 0
# Reorder columns
cols = ['paw_x_pos', 'paw_y_pos', 'paw_x_vel', 'paw_y_vel', 'paw_speed',
'wheel_vel', 'wheel_speed', 'wheel_acc', 'times',
'e_mode', 'e_var', 'frame_id']
data_df = data_df[cols]
return data_df, any_transition_frames
[docs]
def duration_data(data_df, any_transition_frames):
"""
Calculate state durations from processed data.
:param data_df: Processed data DataFrame
:param any_transition_frames: Indices of transition frames
:returns: Tuple of (state_duration_df, duration_summary)
"""
sd_df = data_df[["e_mode", "times", "frame_id"]].copy()
# Filter out noise frames around transitions
for noise_frame in any_transition_frames:
last_persisted_state_idx = max(0, noise_frame - SILENCE_WINDOW - 1)
sd_df.loc[noise_frame - SILENCE_WINDOW:noise_frame + SILENCE_WINDOW, 'e_mode'] = \
sd_df.loc[last_persisted_state_idx, 'e_mode']
# Identify continuous state segments
sd_df['group'] = (sd_df['e_mode'] != sd_df['e_mode'].shift()).cumsum()
# Calculate frame durations
sd_df['frame_duration'] = sd_df['times'].shift(-1) - sd_df['times']
sd_df['frame_duration'] = sd_df['frame_duration'].fillna(1 / 60) # Handle last frame
# Sum durations within each group
durations = sd_df.groupby('group').agg(
e_mode=('e_mode', 'first'),
duration=('frame_duration', 'sum')
).reset_index(drop=True)
return sd_df, durations
[docs]
def interval_data(data_df, trials_data):
"""
Process trial interval data for raster plots.
:param data_df: Processed data DataFrame
:param trials_data: Trials data object
:returns: DataFrame with trial intervals and metadata
"""
ct = data_df["times"]
tr = trials_data
interval_df = tr[['intervals_0', 'intervals_1']].rename(
columns={'intervals_0': 'start', 'intervals_1': 'end'}
)
interval_df["frame_start"] = np.searchsorted(ct, interval_df["start"]).astype(int)
interval_df["frame_end"] = np.searchsorted(ct, interval_df["end"]).astype(int)
interval_df["firstMovement_times"] = tr["firstMovement_times"]
interval_df["frame_fmt"] = np.searchsorted(ct, interval_df["firstMovement_times"])
interval_df["stimOn_times"] = tr["stimOn_times"]
interval_df["feedback_times"] = tr["feedback_times"]
interval_df["trial_idx"] = np.arange(0, len(interval_df))
interval_df["choice_data"] = (tr["feedbackType"] + 1) / 2
interval_df["trial_duration"] = interval_df["end"] - interval_df["start"]
# Calculate average variance per trial
def calculate_avg_var(row):
frame_start = max(0, int(row["frame_start"]))
frame_end = min(len(data_df) - 1, int(row["frame_end"]))
return data_df["e_var"].iloc[frame_start:frame_end].mean()
interval_df["avg_var"] = interval_df.apply(calculate_avg_var, axis=1)
return interval_df
[docs]
def raster_data(interval_df, data_df, fps):
"""
Prepare raster plot data for ensemble modes and variances.
:param interval_df: Trial interval DataFrame
:param data_df: Processed data DataFrame
:param fps: Camera frame rate
:returns: Tuple of (ensemble_raster, variance_raster)
"""
max_frames = (interval_df["frame_end"] - interval_df["frame_start"]).max()
ens_raster_holder = np.full((len(interval_df), max_frames), np.nan)
var_raster_holder = np.full((len(interval_df), max_frames), np.nan)
previous_space = int(fps // 2)
class_to_int = {label: i for i, label in enumerate(STATE_LABELS)}
for idx, row in interval_df.iterrows():
start = int(row["frame_fmt"])
end = int(row["frame_end"])
trial_data = data_df["e_mode"][start - previous_space:end]
trial_length = min(len(trial_data), max_frames)
ens_raster_holder[idx, :trial_length] = trial_data.map(class_to_int).values[:trial_length]
var_raster_holder[idx, :trial_length] = data_df["e_var"][start - previous_space:end].values[:trial_length]
return ens_raster_holder, var_raster_holder
[docs]
def wheel_data(data_df):
"""
Identify transition frames for wheel and movement analysis.
:param data_df: Processed data DataFrame
:returns: Dictionary with transition frame indices
"""
# Add state change tracking
data_df = data_df.copy()
data_df['state_change'] = (data_df[MODEL_TO_USE] != data_df[MODEL_TO_USE].shift()).cumsum()
data_df['state_duration'] = data_df.groupby('state_change').cumcount() + 1
data_df['state_duration_next'] = data_df[::-1].groupby('state_change').cumcount() + 1
# Find transition frames with duration requirements
transitions = {}
transitions['still_to_wheel_turn'] = data_df[
(data_df[MODEL_TO_USE].shift() == 'Still') &
(data_df[MODEL_TO_USE] == 'Wheel Turn') &
(data_df['state_duration'].shift() >= FRAME_COUNTER_DURATION) &
(data_df['state_duration_next'] >= FRAME_COUNTER_DURATION)
].index
transitions['wheel_turn_to_still'] = data_df[
(data_df[MODEL_TO_USE].shift() == 'Wheel Turn') &
(data_df[MODEL_TO_USE] == 'Still') &
(data_df['state_duration'].shift() >= FRAME_COUNTER_DURATION) &
(data_df['state_duration_next'] >= FRAME_COUNTER_DURATION)
].index
transitions['still_to_move'] = data_df[
(data_df[MODEL_TO_USE].shift() == 'Still') &
(data_df[MODEL_TO_USE] == 'Move') &
(data_df['state_duration'].shift() >= FRAME_COUNTER_DURATION) &
(data_df['state_duration_next'] >= FRAME_COUNTER_DURATION)
].index
transitions['move_to_still'] = data_df[
(data_df[MODEL_TO_USE].shift() == 'Move') &
(data_df[MODEL_TO_USE] == 'Still') &
(data_df['state_duration'].shift() >= FRAME_COUNTER_DURATION) &
(data_df['state_duration_next'] >= FRAME_COUNTER_DURATION)
].index
return transitions
# -------------------------------------------------
# Individual plotting functions
# -------------------------------------------------
[docs]
def plot_paw_positions_by_state(ax, frame, data_df, state, state_idx, camera, paw, tracker):
"""
Plot paw positions overlaid on video frame for a specific behavioral state.
:param ax: Matplotlib axis to plot on
:param frame: Video frame as numpy array
:param data_df: Processed data DataFrame
:param state: Behavioral state name
:param state_idx: Index of state for coloring
:param camera: Camera view name
:param paw: Paw identifier
:param tracker: Tracker type
"""
# Filter data for current state
state_data = data_df[data_df['e_mode'] == state]
if len(state_data) == 0:
ax.text(0.5, 0.5, f'No {state} data', ha='center', va='center', transform=ax.transAxes)
ax.axis('off')
return
# Normalize paw speed for visualization
state_data = state_data.copy()
if state_data['paw_speed'].max() > 0:
state_data['paw_speed_normalized'] = (
np.log(state_data['paw_speed'] + 1e-6) / np.log(state_data['paw_speed'].max() + 1e-6)
)
else:
state_data['paw_speed_normalized'] = 0
# Plot frame and overlay points
ax.imshow(frame, cmap='gray')
sc = ax.scatter(
state_data['paw_x_pos'],
state_data['paw_y_pos'],
c=state_data['paw_speed_normalized'],
cmap="plasma",
alpha=0.8,
marker='+',
s=0.1,
vmin=0,
vmax=1
)
ax.set_title(f"Paw Positions: {state}")
ax.axis('off')
# Add colorbar on last panel
if state_idx == len(STATE_LABELS) - 1 and len(state_data) > 0:
cbaxes = inset_axes(ax, width="5%", height="100%", loc='center right', borderpad=0)
cbar_obj = plt.colorbar(sc, cax=cbaxes, orientation='vertical')
cbar_obj.outline.set_visible(True)
cbar_obj.set_ticks([])
cbaxes.text(1.1, 0, 'slow', transform=cbaxes.transAxes,
ha='left', va='bottom', color='black', fontsize=12, rotation=90)
cbaxes.text(1.1, 1, 'fast', transform=cbaxes.transAxes,
ha='left', va='top', color='black', fontsize=12, rotation=90)
[docs]
def plot_state_duration_histogram(ax, durations, state, state_idx, xlim=None):
"""
Plot histogram of durations for a specific behavioral state.
:param ax: Matplotlib axis to plot on
:param durations: DataFrame with state durations
:param state: Behavioral state name
:param state_idx: Index of state for coloring
:param xlim: x-limits of plot if computed externally
"""
state_data = durations[durations['e_mode'] == state]['duration']
if len(state_data) == 0:
ax.text(0.5, 0.5, f'No {state} data', ha='center', va='center', transform=ax.transAxes)
return
sns.histplot(state_data, bins=50, color=COLORS[state_idx], ax=ax,
log_scale=(True, False), fill=True)
ax.set_title(f'Durations: {state}')
ax.set_xlabel('Duration (s)')
if xlim is not None:
ax.set_xlim(xlim)
if state_idx == 0:
ax.set_ylabel('Frequency')
else:
ax.set_ylabel('')
ax.set_yscale('log')
# Set consistent formatting for log scale
ax.yaxis.set_major_locator(ticker.LogLocator(base=10.0))
ax.yaxis.set_minor_locator(ticker.LogLocator(base=10.0, subs=np.arange(2, 10) * 0.1, numticks=10))
ax.yaxis.set_major_formatter(ticker.LogFormatterSciNotation())
ax.yaxis.set_minor_formatter(ticker.NullFormatter())
ax.xaxis.set_major_locator(ticker.LogLocator(base=10.0))
ax.xaxis.set_minor_locator(ticker.LogLocator(base=10.0, subs=np.arange(2, 10) * 0.1, numticks=10))
ax.xaxis.set_major_formatter(ticker.LogFormatterSciNotation())
ax.xaxis.set_minor_formatter(ticker.NullFormatter())
[docs]
def plot_total_duration_bars(ax, durations):
"""
Plot bar chart of total duration percentages for each behavioral state.
:param ax: Matplotlib axis to plot on
:param durations: DataFrame with state durations
"""
# Calculate total duration per state
total_durations = durations.groupby('e_mode')['duration'].sum().reset_index()
# Compute percentage proportion
total_sum = total_durations['duration'].sum()
total_durations['percentage'] = total_durations['duration'] / total_sum * 100
# Map colors
color_map = {state: color for state, color in zip(STATE_LABELS, COLORS)}
total_durations['color'] = total_durations['e_mode'].map(color_map)
cols = total_durations['color'].tolist()
bp = sns.barplot(data=total_durations, x='e_mode', y='percentage',
hue='e_mode', palette=cols, ax=ax, legend=False, order=STATE_LABELS)
# Add value labels on bars
for container in bp.containers:
bp.bar_label(container, padding=3, fmt="%.1f", label_type='center')
ax.set_xlabel('States')
ax.set_ylabel('Total Duration (%)')
ax.set_yscale('log')
[docs]
def plot_ensemble_variance_histogram(ax, data_df):
"""
Plot histogram of ensemble mode variance values.
:param ax: Matplotlib axis to plot on
:param data_df: Processed data DataFrame
"""
sns.histplot(data_df["e_var"], bins=50, color='k', log_scale=False, fill=True, ax=ax)
ax.set_xlabel("Ensemble Mode Variance")
ax.set_xlim(0, 1)
ax.set_ylabel("Frequency")
ax.set_yscale('log')
[docs]
def plot_state_raster(ax, er, fps):
"""
Plot raster of behavioral states across trials.
:param ax: Matplotlib axis to plot on
:param er: Ensemble raster data array
:param fps: Camera frame rate
"""
xlim = 2 * fps # 2 seconds in frames
cmap = ListedColormap(COLORS)
sns.heatmap(er, cmap=cmap, cbar=False, ax=ax, vmin=0, vmax=len(COLORS) - 1)
ax.set_title("Raster of Ensemble Mode States")
ax.set_xlabel("Time from first movement onset (s)")
# Set up y-axis ticks
num_trials = len(er)
max_ticks = 50
if num_trials <= max_ticks:
yt = np.arange(1, num_trials + 1)
else:
yt = np.linspace(1, num_trials, max_ticks, dtype=int)
ax.set_yticks(yt)
ax.set_yticklabels(yt, rotation=0, fontsize=10)
ax.set_ylabel("")
# Set up x-axis
x_ticks = np.linspace(-0.5, 1.5, 9)
ax.set_xlim(0, xlim)
ax.set_xticks(np.linspace(0, xlim, 9))
ax.set_xticklabels([f"{tick:.1f}s" for tick in x_ticks], rotation=0)
# Add vertical line at first movement onset
ax.axvline(x=fps / 2, color='black', linestyle='--')
[docs]
def plot_variance_raster(ax, vr, fps):
"""
Plot raster of ensemble variance values across trials.
:param ax: Matplotlib axis to plot on
:param vr: Variance raster data array
:param fps: Camera frame rate
"""
xlim = 2 * fps # 2 seconds in frames
sns.heatmap(vr, cmap="gray", cbar=False, ax=ax)
ax.set_title("Raster of Ensemble Mode Variance")
ax.set_xlabel("Time from first movement onset (s)")
ax.set_yticks([])
ax.set_xlim(0, xlim)
# Set up x-axis
x_ticks = np.linspace(-0.5, 1.5, 9)
ax.set_xticks(np.linspace(0, xlim, 9))
ax.set_xticklabels([f"{tick:.1f}s" for tick in x_ticks], rotation=0)
# Add vertical line at first movement onset
ax.axvline(x=fps / 2, color='white', linestyle='--')
[docs]
def plot_wheel_speed_transitions(ax, wheel_transitions, data_df, fps, transition_start='still'):
"""
Plot wheel speed around state transitions.
:param ax: Matplotlib axis to plot on
:param wheel_transitions: Dictionary of transition frame indices
:param data_df: Processed data DataFrame
:param fps: Camera frame rate
:param transition_start: which state to start with
"""
# Extract data for transitions
if transition_start == 'still':
still_to_wheel = extract_transition_data(
data_df, wheel_transitions['still_to_wheel_turn'], 'wheel_speed',
)
wheel_to_still = []
elif transition_start == 'wheel_turn':
wheel_to_still = extract_transition_data(
data_df, wheel_transitions['wheel_turn_to_still'], 'wheel_speed',
)
still_to_wheel = []
if len(still_to_wheel) == 0 and len(wheel_to_still) == 0:
ax.text(0.5, 0.5, 'No transition data', ha='center', va='center', transform=ax.transAxes)
return
x = np.arange(-30, 50, 1) / fps # Convert to seconds
alpha = 0.2
# Plot individual traces
for line in still_to_wheel:
ax.plot(x, line, color="gray", alpha=alpha)
for line in wheel_to_still:
ax.plot(x, line, color="gray", alpha=alpha)
# Plot averages
if len(still_to_wheel) > 0:
ax.plot(x, still_to_wheel.mean(axis=0), "k-", label="Still → Wheel Turn", linewidth=2)
if len(wheel_to_still) > 0:
ax.plot(x, wheel_to_still.mean(axis=0), "k-", label="Wheel Turn → Still", linewidth=2)
ax.axvline(0, color='k', linestyle='--', alpha=0.5)
ax.set_xlabel("Time from Transition (seconds)")
ax.set_ylabel("Wheel Speed")
ax.legend(loc='upper right')
ax.grid(True, alpha=0.3)
[docs]
def plot_paw_speed_transitions(ax, wheel_transitions, data_df, fps):
"""
Plot paw speed around various state transitions.
:param ax: Matplotlib axis to plot on
:param wheel_transitions: Dictionary of transition frame indices
:param data_df: Processed data DataFrame
:param fps: Camera frame rate
"""
# Extract data for all transition types
transitions = {}
labels = ["Still → Wheel Turn", "Wheel Turn → Still", "Still → Move", "Move → Still"]
colors = ['b', 'r', 'g', 'orange']
for i, (key, label) in enumerate(zip(
['still_to_wheel_turn', 'wheel_turn_to_still', 'still_to_move', 'move_to_still'],
labels)):
data = extract_transition_data(data_df, wheel_transitions[key], 'paw_speed')
if len(data) > 0:
transitions[key] = data
if len(transitions) == 0:
ax.text(0.5, 0.5, 'No transition data', ha='center', va='center', transform=ax.transAxes)
return
x = np.arange(-30, 50, 1) / fps
# Plot average traces
for i, (key, label, color) in enumerate(zip(
['still_to_wheel_turn', 'wheel_turn_to_still', 'still_to_move', 'move_to_still'],
labels, colors)):
if key in transitions:
ax.plot(x, transitions[key].mean(axis=0), color=color, label=label, linewidth=2)
ax.axvline(0, color='k', linestyle='--', label="State Transition")
ax.set_xlabel("Time from Transition (s)")
ax.set_ylabel("Paw Speed")
ax.legend()
ax.grid(True, alpha=0.3)
[docs]
def plot_trial_correctness(ax, interval_df):
"""
Plot trial correctness as a binary heatmap.
:param ax: Matplotlib axis to plot on
:param interval_df: Trial interval DataFrame
"""
binary_cmap = ListedColormap(['black', 'white'])
binary_data = np.expand_dims(interval_df['choice_data'].values, axis=1)
sns.heatmap(binary_data, cmap=binary_cmap, cbar=False, ax=ax)
ax.set_xticks([])
ax.set_yticks([])
ax.set_ylabel("Incorrect Response", fontsize=12)
[docs]
def plot_trial_duration(ax, interval_df):
"""
Plot trial duration as a grayscale heatmap.
:param ax: Matplotlib axis to plot on
:param interval_df: Trial interval DataFrame
"""
duration_data = np.expand_dims(interval_df['trial_duration'].values, axis=1)
sns.heatmap(duration_data, cmap="gray", cbar=False, ax=ax)
ax.set_xticks([])
ax.set_yticks([])
ax.set_ylabel("Trial Duration", fontsize=12)