How to work with brain states and regression

cogpy.brainstates provides interval-based state classification (sleep stages, behavioral states, transitions). cogpy.regression builds design matrices and solves OLS problems for artifact removal or signal decomposition. They are independent modules that compose naturally in state-stratified analyses.

State dictionary format

Brain states are represented as a dict mapping labels to lists of [start, end) intervals (seconds):

states = {
    "PerSWS": [[10.0, 45.0], [120.0, 200.0]],
    "PerREM": [[50.0, 80.0], [210.0, 260.0]],
}

Convert states to a DataFrame

from cogpy.brainstates.brainstates import get_states_df

states_df = get_states_df(states)
# Columns: state, iseg, t0, t1

Get per-state durations

from cogpy.brainstates.brainstates import get_state_durations

durations = get_state_durations(states_df)
# Returns pd.Series indexed by state label

Label time points by state

Given a time array and a state dict, assign each sample to a state period:

from cogpy.brainstates.brainstates import label_numbers_by_state_intervals

t = sig.coords["time"].values
labeled = label_numbers_by_state_intervals(t, states)
# DataFrame: one column per state, values = period index or -1

Filter events by state

from cogpy.brainstates.brainstates import sort_col_into_states, filter_by_states

# Add state labels to an event DataFrame
events_df = sort_col_into_states(events_df, "time", states)

# Keep only events during SWS
sws_events = filter_by_states(events_df, include_states=["PerSWS"], exclude_states=[])

Restrict signal to intervals

from cogpy.brainstates.intervals import restrict

# Keep only SWS portions of the signal
sig_sws = restrict(sig, states["PerSWS"])

# Or pass the full states dict (union of all intervals)
sig_any_state = restrict(sig, states)

Detect state transitions

from cogpy.brainstates.brainstates import state_transitions, state_transition_interval

# All transitions
trans_df = state_transitions(states)
# Columns: transition_time, prev_state, next_state

# Extract SWS→REM transition windows (30 s before, 30 s after)
windows = state_transition_interval(states, "PerSWS", "PerREM", 30.0, 30.0)
# Returns (n_transitions, 2) array of [t_start, t_end]

Add transition intervals back as a new state:

from cogpy.brainstates.brainstates import append_transition_intervals

states = append_transition_intervals(states, "PerSWS", "PerREM", window_before=30, window_after=0)
# Adds key "PerSWS_To_PerREM" to states dict

Purify macro states (remove micro-state overlap)

When macrostates (e.g. SWS) contain embedded microstates (e.g. high-velocity spindles), subtract the microstates:

from cogpy.brainstates.brainstates import purify_macro_states

states = purify_macro_states(
    states,
    macro_states=["PerSWS", "PerREM"],
    micro_states=["PerHVS", "PerMicroA"],
)

EMG proxy from intracranial recordings

Estimate muscle activity from high-frequency inter-channel correlations (Watson & Buzsaki protocol):

from cogpy.brainstates.EMG import compute_emg_proxy

emg_df = compute_emg_proxy(
    sig,
    fs=1000.0,
    coords_df=channel_coords_df,   # columns: ml, ap, dv (micrometers)
    min_distance=400.0,             # only use distant channel pairs
    window_size=0.5,                # 500 ms windows
    window_step=0.25,               # 250 ms step
)
# Returns DataFrame with columns: time, emg_proxy, emg_proxy_std

Build a lagged design matrix

lagged_design_matrix creates a Toeplitz matrix where each column is the reference signal shifted by a given lag. Positive lags look backward in time.

from cogpy.regression import lagged_design_matrix

# Reference signal (e.g. stimulation TTL), lags 0..5 samples
X = lagged_design_matrix(reference, lags=range(6), intercept=True)
# X.shape → (n_time, 7)  — 1 intercept + 6 lag columns

Build an event design matrix

When you have discrete events and a known template waveform, build a design matrix with one column per event:

from cogpy.regression import event_design_matrix

# template: (n_lag,) waveform, event_samples: (n_events,) sample indices
X = event_design_matrix(len(sig), event_samples, template, intercept=True)
# X.shape → (n_time, n_events + 1)

Fit, predict, and clean with OLS

from cogpy.regression import ols_fit, ols_predict, ols_residual

# Y: (n_time,) or (n_time, n_channels) — supports multi-channel
beta = ols_fit(X, Y)            # least-squares coefficients
Y_hat = ols_predict(X, beta)    # predicted artifact
cleaned = ols_residual(X, Y, beta)  # Y - Y_hat

Full example: state-stratified artifact removal

Remove stimulation artifacts from SWS epochs using lagged regression:

from cogpy.brainstates.intervals import restrict
from cogpy.regression import lagged_design_matrix, ols_fit, ols_residual
import numpy as np

# 1. Restrict signal and reference to SWS intervals
sig_sws = restrict(sig, states["PerSWS"])
ref_sws = restrict(stim_ref, states["PerSWS"])

# 2. Build lagged design matrix (0–10 ms at 1 kHz = lags 0..10)
X = lagged_design_matrix(ref_sws.values, lags=range(11), intercept=True)

# 3. Fit and subtract per-channel
Y = sig_sws.values  # (n_time, n_channels)
beta = ols_fit(X, Y)
sig_clean = ols_residual(X, Y, beta)

Full example: event-locked template regression

Remove a known artifact template at each event:

from cogpy.triggered import estimate_template
from cogpy.brainstates.intervals import perievent_epochs
from cogpy.regression import event_design_matrix, ols_fit, ols_residual
import numpy as np

# 1. Estimate template from epochs
epochs = perievent_epochs(sig, event_times, fs=1000.0, pre=0.001, post=0.010)
template = estimate_template(epochs, method="median").values.ravel()

# 2. Build design matrix
event_samples = np.round(event_times * 1000).astype(int)
X = event_design_matrix(sig.sizes["time"], event_samples, template)

# 3. Fit per-event amplitudes and subtract
Y = sig.values
beta = ols_fit(X, Y)
cleaned = ols_residual(X, Y, beta)