"""Public pipeline module. Conceptually, this module converts one or more NWB sessions into canonical PSTH entries enriched with behavioral traces, quiet-trial masks, and unit metadata. It exists as a separate unit so NWB I/O, trial-aligned signal construction, and export-ready entry assembly can stay coordinated without being repeated in notebooks. It connects NWB session files to the canonical per-area entry structures used throughout the analysis pipeline. """ from __future__ import annotations import pickle import zipfile from dataclasses import dataclass from pathlib import Path from typing import Any, Dict, List, Mapping, Sequence import numpy as np import pandas as pd from pynwb import NWBHDF5IO from scipy.signal import savgol_filter from tqdm.auto import tqdm from .behavior import QuietTrialParams, nwb_find_quiet_trial, psth_behavior, psth_simple, robust_mode, trial_type_maker from .loading import ( get_nwb_trials_dataframe, load_nwb_behavior_timeseries_data, load_nwb_behavior_timeseries_timestamps, make_picklable, matlab_like_trial_columns, ) @dataclass(frozen=True) class PsthBuildConfig: """Configuration bundle for canonical PSTH-session construction.""" pre_time: float post_time: float spike_bin_width: float spike_bin_step: float quiet_params: QuietTrialParams piezo_bin_width: float | None = None piezo_bin_step: float | None = None movement_bin_width: float = 0.01 movement_bin_step: float = 0.01 piezo_sampling_rate: float = 2000.0 pg019_frame_rate: float = 100.0 default_frame_rate: float = 200.0 # MATLAB keeps all NWB timestamps in double precision. Trial start times are # large absolute session timestamps (thousands of seconds); casting the # spike/behaviour anchors to float32 (~7 significant digits) drops # sub-millisecond precision and shifts spikes across 10 ms bin edges relative # to Code_M, which propagates into divergent clustering. Keep float64. timestamp_dtype: Any = np.float64 trial_timestamps_source: str = "behavior" include_behaviour_timestamps: bool = False extra_trial_fields: tuple[str, ...] = ("lick_time", "early_lick") @property def resolved_piezo_bin_width(self) -> float: return self.spike_bin_width if self.piezo_bin_width is None else float(self.piezo_bin_width) @property def resolved_piezo_bin_step(self) -> float: return self.spike_bin_step if self.piezo_bin_step is None else float(self.piezo_bin_step) def _smooth_piezo_trace(values: np.ndarray) -> np.ndarray: """Return a lightly smoothed piezo trace while preserving short signals.""" values = np.asarray(values, dtype=np.float32).reshape(-1) if values.size >= 11: return savgol_filter(values, window_length=11, polyorder=5, mode="interp").astype(np.float32) return values.astype(np.float32, copy=False) def _compute_movement_signal( jaw_coord: np.ndarray, tongue_coord: np.ndarray, snout_angle: np.ndarray, whisker_angle: np.ndarray, ) -> Dict[str, np.ndarray]: """Derive the movement signals used by the behavioral PSTH builder.""" if whisker_angle.size <= 1: whisker_speed = np.zeros_like(whisker_angle, dtype=np.float32) else: diff_abs = np.abs(np.diff(whisker_angle.astype(np.float32), n=1)) whisker_speed = np.concatenate([diff_abs[:1], diff_abs]).astype(np.float32) jaw_coord = np.asarray(jaw_coord, dtype=np.float32) tongue_coord = np.asarray(tongue_coord, dtype=np.float32) jaw_x = jaw_coord[:, 0] jaw_y = jaw_coord[:, 1] jaw_mode_x = float(robust_mode(jaw_x)) jaw_mode_y = float(robust_mode(jaw_y)) jaw_movement = np.sqrt((jaw_y - jaw_mode_y) ** 2 + (jaw_x - jaw_mode_x) ** 2).astype(np.float32) tongue_x = tongue_coord[:, 0] tongue_y = tongue_coord[:, 1] tongue_movement = np.sqrt((tongue_y - jaw_mode_y) ** 2 + (tongue_x - jaw_mode_x) ** 2).astype(np.float32) tongue_mean = float(np.nanmean(tongue_movement)) if tongue_movement.size else np.nan if np.isfinite(tongue_mean) and tongue_mean > 0: # Drop extreme tongue excursions so one tracking glitch does not # dominate the trial-aligned averages. tongue_movement[tongue_movement > (20.0 * tongue_mean)] = np.nan return { "whisker_angle": np.asarray(whisker_angle, dtype=np.float32), "whisker_speed": whisker_speed.astype(np.float32, copy=False), "jaw_movement": jaw_movement, "snout_angle": np.asarray(snout_angle, dtype=np.float32), "tongue_movement": tongue_movement, } def _build_behavior_table( nwb: Any, anchor_times: np.ndarray, nwb_path: Path, cfg: PsthBuildConfig, ) -> tuple[Dict[str, Any], Dict[str, np.ndarray], Dict[str, np.ndarray]]: """Build trial-aligned behavioral payloads for one NWB session. Parameters ---------- nwb: Open NWB object for the current session. anchor_times: Trial start times used to align behavioral traces. nwb_path: Path to the current NWB file, used for session-specific sampling-rate adjustments. cfg: PSTH build configuration bundle. Returns ------- tuple[dict[str, Any], dict[str, np.ndarray], dict[str, np.ndarray]] Behavior table ready to merge into each area entry, plus the movement and piezo window-center mappings used during alignment. """ jaw_coord = np.asarray(load_nwb_behavior_timeseries_data(nwb, "BehavioralTimeSeries", "Jaw_Coordinate")) tongue_coord = np.asarray(load_nwb_behavior_timeseries_data(nwb, "BehavioralTimeSeries", "Tongue_Coordinate")) snout_angle = np.asarray( load_nwb_behavior_timeseries_data(nwb, "BehavioralTimeSeries", "Snout_Angle") ).reshape(-1) whisker_angle = np.asarray( load_nwb_behavior_timeseries_data(nwb, "BehavioralTimeSeries", "C2Whisker_Angle") ).reshape(-1) video_timestamp = np.asarray( load_nwb_behavior_timeseries_timestamps(nwb, "BehavioralTimeSeries", "C2Whisker_Angle"), dtype=cfg.timestamp_dtype, ) piezo_lick_trace_raw = np.asarray( load_nwb_behavior_timeseries_data(nwb, "BehavioralTimeSeries", "Piezo_lick_trace") ).reshape(-1) piezo_timestamp = np.asarray( load_nwb_behavior_timeseries_timestamps(nwb, "BehavioralTimeSeries", "Piezo_lick_trace"), dtype=cfg.timestamp_dtype, ) movement_signal = _compute_movement_signal(jaw_coord, tongue_coord, snout_angle, whisker_angle) piezo_lick_trace = _smooth_piezo_trace(piezo_lick_trace_raw) # Some recordings use a different camera rate, so resolve the frame rate # before building movement PSTHs. frame_rate = cfg.pg019_frame_rate if "PG019" in nwb_path.name else cfg.default_frame_rate psth_beh, window_centers = psth_behavior( movement_signal, video_timestamp, anchor_times, cfg.pre_time, cfg.post_time, cfg.movement_bin_width, cfg.movement_bin_step, frame_rate, ) psth_piezo, piezo_window_centers = psth_behavior( {"piezo_lick_trace": piezo_lick_trace}, piezo_timestamp, anchor_times, cfg.pre_time, cfg.post_time, cfg.resolved_piezo_bin_width, cfg.resolved_piezo_bin_step, cfg.piezo_sampling_rate, ) behavior_table: Dict[str, Any] = dict(psth_beh) behavior_table["piezo_lick_trace"] = np.asarray(psth_piezo["piezo_lick_trace"], dtype=np.float32) behavior_table["trial_timestamps"] = np.asarray(window_centers["whisker_angle"], dtype=np.float32) behavior_table = nwb_find_quiet_trial(behavior_table, cfg.quiet_params) return behavior_table, window_centers, piezo_window_centers def _default_spike_window_centers(cfg: PsthBuildConfig) -> np.ndarray: """Return the default spike-window centers implied by the PSTH config.""" edges = np.arange(cfg.pre_time, cfg.post_time + 1e-12, cfg.spike_bin_step, dtype=float) n_bins = max(int(edges.size - 1), 0) return np.asarray(edges[:n_bins] + float(cfg.spike_bin_width), dtype=np.float32) def _build_area_entry( units_area: pd.DataFrame, elec_area: pd.DataFrame, trials_df: pd.DataFrame, behavior_table: Mapping[str, Any], anchor_times: np.ndarray, piezo_window_centers: Mapping[str, np.ndarray], session_id: str, area_name: str, cfg: PsthBuildConfig, nwb_path: Path, ) -> Dict[str, Any]: """Assemble one canonical PSTH entry for a single session-area pair. Parameters ---------- units_area, elec_area: Area-restricted unit and electrode tables. trials_df: Session trial table. behavior_table: Trial-aligned behavioral payload returned by `_build_behavior_table`. anchor_times: Trial start times used to align spike counts. piezo_window_centers: Window-center mapping returned by `_build_behavior_table`. session_id: Session identifier stored in the output entry. area_name: Area label for the current entry. cfg: PSTH build configuration bundle. nwb_path: Source NWB path, used only for informative error messages. Returns ------- dict[str, Any] One canonical per-area entry ready to save into processed data. """ entry: Dict[str, Any] = { "session_id": session_id, "probe_location": area_name, } # Code_M uses `fieldnames(units_table)(1:end-3)`, but for a MATLAB table the # last 3 fieldnames are the pseudo-fields (Properties/Row/Variables), not real # data columns. pandas DataFrames have no such pseudo-fields, so dropping 3 # real columns here deletes genuine unit metadata (e.g. # `unit_allenccf_area_layer` -> blank Layer panel). Keep every real column and # strip only the nested `electrodes` column below. unit_cols = matlab_like_trial_columns(units_area.columns, drop_first=0, drop_last=0) if not unit_cols: unit_cols = [c for c in units_area.columns if c != "electrodes"] unit_cols = [c for c in unit_cols if c != "electrodes"] for col in unit_cols: entry[f"unit_{col}"] = make_picklable(units_area[col].to_numpy()) # Same rationale as the unit columns: keep all real electrode columns (the # MATLAB `(1:end-3)` trim only removes table pseudo-fields) and drop only the # nested `group` column below. elec_cols = matlab_like_trial_columns(elec_area.columns, drop_first=0, drop_last=0) if not elec_cols: elec_cols = [c for c in elec_area.columns if c != "group"] elec_cols = [c for c in elec_cols if c != "group"] for col in elec_cols: entry[f"elec_{col}"] = make_picklable(elec_area[col].to_numpy()) entry.pop("unit_electrodes", None) entry.pop("elec_group", None) # Copy the shared trial-aligned behavior payload into every area entry so # downstream code can work area-by-area without a separate session lookup. for key, value in behavior_table.items(): entry[key] = make_picklable(value) # Same rationale: keep all real trial columns; the MATLAB `(1:end-3)` trim # only removes table pseudo-fields, which pandas does not have. trial_cols = matlab_like_trial_columns(trials_df.columns, drop_first=0, drop_last=0) if not trial_cols: trial_cols = list(trials_df.columns) for col in trial_cols: entry[col] = make_picklable(trials_df[col].to_numpy()) for extra_trial_col in cfg.extra_trial_fields: if extra_trial_col in trials_df.columns and extra_trial_col not in entry: entry[extra_trial_col] = make_picklable(trials_df[extra_trial_col].to_numpy()) if {"whisker_stim", "context"}.issubset(trials_df.columns): entry["trial_type"] = trial_type_maker(trials_df["whisker_stim"], trials_df["context"]).astype(np.int16) else: entry["trial_type"] = np.zeros(trials_df.shape[0], dtype=np.int16) if "spike_times" not in units_area.columns: raise KeyError(f"units table missing 'spike_times': {nwb_path.name} | {area_name}") spike_times_list = list(units_area["spike_times"].to_numpy()) spike_counts_all: List[np.ndarray] = [] spike_window_centers = _default_spike_window_centers(cfg) # Build one `(bins, trials)` spike-count matrix per unit, then stack units # along the last axis to obtain the canonical `(bins, trials, cells)` tensor. for spike_times in spike_times_list: _, curr_window_centers, spike_counts = psth_simple( np.asarray(spike_times, dtype=float), np.asarray(anchor_times, dtype=float), cfg.pre_time, cfg.post_time, cfg.spike_bin_width, cfg.spike_bin_step, ) spike_window_centers = np.asarray(curr_window_centers, dtype=np.float32) spike_counts_all.append(np.asarray(spike_counts, dtype=np.float32)) if spike_counts_all: entry["spike_counts"] = np.stack(spike_counts_all, axis=2).astype(np.float32) else: entry["spike_counts"] = np.zeros((spike_window_centers.size, anchor_times.size, 0), dtype=np.float32) if cfg.trial_timestamps_source == "spike": entry["trial_timestamps"] = spike_window_centers.astype(np.float32) elif cfg.trial_timestamps_source != "behavior": raise ValueError(f"Unknown trial_timestamps_source: {cfg.trial_timestamps_source}") if cfg.include_behaviour_timestamps: entry["behaviour_timestamps"] = np.asarray( piezo_window_centers["piezo_lick_trace"], dtype=np.float32, ) return entry def build_ephys_psth_session_from_nwb_path( nwb_path: Path, cfg: PsthBuildConfig, ) -> Dict[str, Any]: """Build all canonical PSTH entries for one NWB session. Parameters ---------- nwb_path: Path to the source NWB file. cfg: PSTH build configuration bundle. Returns ------- dict[str, Any] Session bundle containing the source path, session identifier, and the list of per-area PSTH entries. """ nwb_path = Path(nwb_path) with NWBHDF5IO(str(nwb_path), mode="r", load_namespaces=True) as io: nwb = io.read() units_df = nwb.units.to_dataframe() if nwb.units is not None else pd.DataFrame() electrodes_df = nwb.electrodes.to_dataframe() if nwb.electrodes is not None else pd.DataFrame() trials_df = get_nwb_trials_dataframe(nwb) if "start_time" not in trials_df.columns: raise KeyError(f"Missing 'start_time' in trials table: {nwb_path.name}") if "location" not in units_df.columns: raise KeyError(f"units table missing 'location': {nwb_path.name}") anchor_times = trials_df["start_time"].to_numpy(dtype=cfg.timestamp_dtype) behavior_table, _, piezo_window_centers = _build_behavior_table(nwb, anchor_times, nwb_path, cfg) # Code_M iterates areas via `unique(units_table.location)`, which returns # them in SORTED (alphabetical) order. GlobalclusterID is then assigned in # that order, so match it with sorted() here — using pandas' occurrence # order (pd.unique) would give multi-area sessions different global unit # ids than Code_M (the neural data is identical, only the Unit_ids labels # shift). area_names = sorted(str(x) for x in pd.unique(units_df["location"].astype(str))) session_id = str(getattr(nwb, "identifier", nwb_path.stem)) session_entries: List[Dict[str, Any]] = [] for area_name in area_names: ind_units = units_df["location"].astype(str).to_numpy() == area_name units_area = units_df.loc[ind_units] if "location" in electrodes_df.columns: ind_elec = electrodes_df["location"].astype(str).to_numpy() == area_name elec_area = electrodes_df.loc[ind_elec] else: elec_area = pd.DataFrame() session_entries.append( _build_area_entry( units_area=units_area, elec_area=elec_area, trials_df=trials_df, behavior_table=behavior_table, anchor_times=anchor_times, piezo_window_centers=piezo_window_centers, session_id=session_id, area_name=area_name, cfg=cfg, nwb_path=nwb_path, ) ) return { "nwb_path": nwb_path, "session_id": session_id, "psth_entries": session_entries, } def build_ephys_psth_sessions( nwb_paths: Sequence[Path], cfg: PsthBuildConfig, show_progress: bool = True, ) -> List[Dict[str, Any]]: """Build canonical PSTH session bundles for multiple NWB files.""" sessions: List[Dict[str, Any]] = [] iter_files = tqdm(nwb_paths, desc="Process NWB", unit="file") if show_progress else nwb_paths for nwb_path in iter_files: sessions.append(build_ephys_psth_session_from_nwb_path(Path(nwb_path), cfg)) return sessions def flatten_psth_sessions(sessions: Sequence[Mapping[str, Any]]) -> List[Dict[str, Any]]: """Flatten session bundles into one list of canonical PSTH entries.""" psth_mat: List[Dict[str, Any]] = [] for sess in sessions: psth_mat.extend(list(sess.get("psth_entries", []))) return psth_mat def assign_global_cluster_ids(entries: Sequence[Dict[str, Any]]) -> int: """Assign monotonically increasing global unit ids across all entries. Parameters ---------- entries: Canonical PSTH entries updated in place. Returns ------- int Total number of units assigned across all entries. """ offset = 0 for entry in entries: n_units = 0 if "unit_cluster_id" in entry: n_units = len(entry["unit_cluster_id"]) elif "unit_id" in entry: n_units = len(entry["unit_id"]) else: spike_counts = np.asarray(entry.get("spike_counts", np.zeros((0, 0, 0), dtype=np.float32))) if spike_counts.ndim == 3: n_units = int(spike_counts.shape[2]) entry["GlobalclusterID"] = np.arange(1, n_units + 1, dtype=np.int32) + offset offset += n_units return offset def save_psth_entries_npz( out_path: Path, entries: Sequence[Mapping[str, Any]], key: str = "psth_mat", *, compressed: bool = False, ) -> None: """Save canonical PSTH entries to an NPZ payload. Large PSTH bundles are saved uncompressed by default to avoid the high peak-memory cost of compressing a pickled object array. """ out_path = Path(out_path) out_path.parent.mkdir(parents=True, exist_ok=True) tmp_path = out_path.with_name(f".{out_path.stem}.tmp{out_path.suffix}") entries_list = list(entries) obj = np.empty(len(entries_list), dtype=object) obj[:] = entries_list if obj.dtype.hasobject and not compressed: try: _save_object_array_npz(tmp_path, key, obj) tmp_path.replace(out_path) except Exception: try: tmp_path.unlink(missing_ok=True) except PermissionError: pass raise return save_fn = np.savez_compressed if compressed else np.savez try: save_fn(tmp_path, **{key: obj}) tmp_path.replace(out_path) except Exception: try: tmp_path.unlink(missing_ok=True) except PermissionError: pass raise def _save_object_array_npz(out_path: Path, key: str, array: np.ndarray) -> None: """Write an object-array NPZ using a newer pickle protocol. NumPy's object-array `.npz` writer currently pickles with protocol 4, which can create large transient allocations for PSTH bundles. Writing the `.npy` member directly keeps the file readable by `np.load(..., allow_pickle=True)`. """ class _PickleZipWriter: """Adapt zipfile handles for protocol-5 PickleBuffer writes.""" def __init__(self, handle: Any) -> None: self._handle = handle def write(self, data: Any) -> int: if isinstance(data, pickle.PickleBuffer): data = data.raw() return self._handle.write(data) header = np.lib.format.header_data_from_array_1_0(array) with zipfile.ZipFile(out_path, mode="w", compression=zipfile.ZIP_STORED, allowZip64=True) as zip_file: with zip_file.open(f"{key}.npy", mode="w", force_zip64=True) as handle: np.lib.format.write_array_header_2_0(handle, header) pickle.dump(array, _PickleZipWriter(handle), protocol=pickle.HIGHEST_PROTOCOL)