"""Public helper module. Conceptually, this module loads processed analysis files and turns their raw payloads into stable canonical structures for downstream decoding workflows. It exists as a separate unit so file-format handling, canonicalization, and cross-file alignment stay centralized instead of being reimplemented in each notebook or pipeline. It connects low-level loading helpers to decoding, spontaneous-lick, and movement-subspace analyses. """ from __future__ import annotations import pickle from pathlib import Path from typing import Any, Dict, List, Optional, Set, Tuple import numpy as np from .constants import FUNCTION_API_VERSION from .core import clean_str, normalize_raw_entries, to_1d, to_1d_str, to_session_id_str, unwrap_scalar_obj from .loading_base import ( _entry_get_any, _normalize_vector_length, sanitize_time_axis, ) __all__ = [ "canonicalize_psth_entries_for_decoding", "decode_area_list", "extract_strings_any", "infer_placeholder_id_order", "load_area_list_from_npz", "load_area_list_from_npz_or_mat", "load_canonical_movement_subspace_entries", "load_canonical_psth_entries_and_area_list", "load_canonical_spontlick_entries", "load_decoding_cluster_inputs", "load_psth_entries_and_area_list", "load_raw_psth_entries", ] def decode_area_list(area_list_obj: Any) -> Optional[Dict[str, Set[str]]]: """Decode an `Area_list` payload into a Python mapping. Parameters ---------- area_list_obj: Raw payload loaded from serialized area-list structures. Returns ------- dict[str, set[str]] or None Decoded area-to-label mapping, or `None` when the payload cannot be interpreted as a flat area list. """ if area_list_obj is None: return None obj = unwrap_scalar_obj(area_list_obj) fields: Dict[str, Any] = {} if isinstance(obj, dict): fields = {str(k): v for k, v in obj.items()} elif hasattr(obj, "_fieldnames"): for field in obj._fieldnames or []: fields[str(field)] = getattr(obj, field) elif isinstance(obj, np.ndarray) and obj.size == 1 and obj.dtype.names: rec = obj.reshape(-1)[0] fields = {str(field): rec[field] for field in rec.dtype.names} decoded: Dict[str, Set[str]] = {} for field, raw in fields.items(): arr_raw = np.asarray(raw) if arr_raw.dtype.names: return None vals: List[str] = [] # Drop stringified empty values early so the returned area map can be # used directly for unit-location membership checks. for value in arr_raw.reshape(-1): if isinstance(value, bytes): s = value.decode("utf-8", errors="ignore").strip() else: s = str(value).strip() if s and s.lower() not in ("none", "nan"): vals.append(s) if vals: decoded[str(field)] = set(vals) return decoded if decoded else None def extract_strings_any(raw: Any) -> List[str]: """Recursively collect clean strings from nested object payloads. Parameters ---------- raw: Arbitrary nested payload loaded from serialized structures. Returns ------- list[str] Flattened list of non-empty strings with placeholder values removed. """ out: List[str] = [] def _rec(v: Any) -> None: vv = unwrap_scalar_obj(v) if vv is None: return if isinstance(vv, (str, bytes)): s = clean_str(vv.decode("utf-8", errors="ignore") if isinstance(vv, bytes) else vv) if s and s.lower() not in ("nan", "none"): out.append(s) return if isinstance(vv, np.ndarray): if vv.dtype.names and vv.size == 1: rec = vv.reshape(-1)[0] for name in vv.dtype.names: _rec(rec[name]) return for value in vv.reshape(-1): _rec(value) return if hasattr(vv, "_fieldnames"): for field in (vv._fieldnames or []): _rec(getattr(vv, field)) return s = clean_str(vv) if s and s.lower() not in ("nan", "none"): out.append(s) _rec(raw) return out def load_area_list_from_npz( npz_path: Optional[Path], ) -> Dict[str, Set[str]]: """Load an `Area_list` mapping from the canonical NPZ payload. Parameters ---------- npz_path: Optional NPZ path containing `area_list` or `Area_list`. Returns ------- dict[str, set[str]] Decoded area-to-CCF mapping. """ npz_path = Path(npz_path) if npz_path is not None else None if npz_path is not None and npz_path.exists(): area_list_obj: Any = None with np.load(npz_path, allow_pickle=True) as data: if "area_list" in data: area_list_obj = data["area_list"] elif "Area_list" in data: area_list_obj = data["Area_list"] decoded_npz = decode_area_list(area_list_obj) if decoded_npz: return decoded_npz raise ValueError("Could not decode Area_list from NPZ") def load_area_list_from_npz_or_mat( npz_path: Optional[Path], mat_path: Optional[Path] = None, ) -> Dict[str, Set[str]]: """Backward-compatible NPZ-only wrapper for legacy imports.""" del mat_path return load_area_list_from_npz(npz_path) def load_psth_entries_and_area_list( psth_path: Path, area_list_path: Path, ) -> Tuple[List[Dict[str, Any]], Dict[str, Set[str]]]: """Load raw PSTH entries and decode the area-list mapping. Parameters ---------- psth_path: NPZ file containing `psth_mat`. area_list_path: NPZ file containing `area_list` or `Area_list`. Returns ------- list[dict[str, Any]], dict[str, set[str]] Raw PSTH entries and the decoded area-list mapping. """ psth_path = Path(psth_path) area_list_path = Path(area_list_path) if not psth_path.exists(): raise FileNotFoundError(f"File not found: {psth_path}") if not area_list_path.exists(): raise FileNotFoundError(f"File not found: {area_list_path}") with np.load(psth_path, allow_pickle=True) as data: if "psth_mat" not in data: raise KeyError("NPZ missing key 'psth_mat'") raw_arr = np.asarray(data["psth_mat"]).reshape(-1) entries_raw = normalize_raw_entries(raw_arr) area_list_obj: Any = None with np.load(area_list_path, allow_pickle=True) as data: if "area_list" in data: area_list_obj = data["area_list"] elif "Area_list" in data: area_list_obj = data["Area_list"] area_list = decode_area_list(area_list_obj) if area_list is None: raise RuntimeError("Area_list decode failed. Fix Area_list.npz decoding first.") return entries_raw, area_list def load_raw_psth_entries(psth_path: Path) -> List[Dict[str, Any]]: """Load raw PSTH entries from one NPZ file without canonicalization. Parameters ---------- psth_path: NPZ file containing the top-level `psth_mat` payload. Returns ------- list[dict[str, Any]] Raw PSTH entries normalized into plain Python mappings. """ psth_path = Path(psth_path) if not psth_path.exists(): raise FileNotFoundError(f"File not found: {psth_path}") with np.load(psth_path, allow_pickle=True) as data: if "psth_mat" not in data: raise KeyError("NPZ missing key 'psth_mat'") raw_arr = np.asarray(data["psth_mat"]).reshape(-1) return normalize_raw_entries(raw_arr) def load_canonical_psth_entries_and_area_list( psth_path: Path, area_list_path: Path, ) -> Dict[str, Any]: """Load PSTH entries and return a canonical decoding-friendly structure. Parameters ---------- psth_path: NPZ file containing `psth_mat`. area_list_path: NPZ file containing the area-list payload. Returns ------- dict[str, Any] Canonical bundle containing cleaned entries, entries grouped by area, canonical window centers, and the decoded area list. """ entries_raw, area_list = load_psth_entries_and_area_list(psth_path, area_list_path) id_ordered, neuron_counter = infer_placeholder_id_order(entries_raw) canon = canonicalize_psth_entries_for_decoding( entries_raw=entries_raw, id_ordered=id_ordered, neuron_counter=neuron_counter, ) return { "entries_raw": entries_raw, "entries_clean": canon["entries_clean"], "entries_by_area": canon["entries_by_area"], "window_centers": canon["window_centers"], "area_list": area_list, "id_ordered": id_ordered, "neuron_counter": neuron_counter, } def infer_placeholder_id_order( entries_raw: List[Dict[str, Any]], ) -> Tuple[np.ndarray, np.ndarray]: """Infer placeholder ID arrays for canonicalization. Parameters ---------- entries_raw: Raw PSTH entries before canonicalization. Returns ------- np.ndarray, np.ndarray Placeholder `id_ordered` and `neuron_counter` vectors. """ gid_max = 0 for entry in entries_raw: raw_gid = entry.get("GlobalclusterID", None) if raw_gid is None: continue gid = to_1d(unwrap_scalar_obj(raw_gid), np.int64) if gid.size: gid_max = max(gid_max, int(np.nanmax(gid))) if gid_max <= 0: total_units = 0 for entry in entries_raw: spike_counts = np.asarray(unwrap_scalar_obj(entry["spike_counts"])) if spike_counts.ndim == 3: total_units += int(spike_counts.shape[2]) gid_max = max(total_units, 1) id_ordered = np.arange(1, gid_max + 1, dtype=np.int64) neuron_counter = id_ordered.copy() return id_ordered, neuron_counter def load_decoding_cluster_inputs( psth_path: Path, area_list_path: Optional[Path], cluster_path: Path, apply_ccf_filter: bool = True, require_area_list: bool = True, require_id_ordered: bool = True, ) -> Dict[str, Any]: """Load and validate inputs for cluster-dropout decoding workflows. Parameters ---------- psth_path: NPZ file containing `psth_mat`. area_list_path: Optional NPZ file containing `area_list` metadata. cluster_path: NPZ or PKL file containing clustering outputs. apply_ccf_filter: Whether decoded area metadata should enable downstream CCF filtering. require_area_list: If True, fail when the area-list payload cannot be decoded. require_id_ordered: If True, fail when the clustering payload does not expose `Id_ordered`. Returns ------- dict[str, Any] Validated inputs ready for cluster-dropout decoding notebooks. """ psth_path = Path(psth_path) cluster_path = Path(cluster_path) area_list_path = Path(area_list_path) if area_list_path is not None else None if not psth_path.exists(): raise FileNotFoundError(f"File not found: {psth_path}") if not cluster_path.exists(): raise FileNotFoundError(f"File not found: {cluster_path}") with np.load(psth_path, allow_pickle=True) as data: if "psth_mat" not in data: raise KeyError("NPZ missing key 'psth_mat'") raw_arr = np.asarray(data["psth_mat"]).reshape(-1) entries_raw = normalize_raw_entries(raw_arr) area_list_obj = None if area_list_path is not None and area_list_path.exists(): with np.load(area_list_path, allow_pickle=True) as data: if "area_list" in data: area_list_obj = data["area_list"] elif "Area_list" in data: area_list_obj = data["Area_list"] area_list = decode_area_list(area_list_obj) enable_ccf_filter = bool(apply_ccf_filter and (area_list is not None)) if require_area_list and area_list is None: raise RuntimeError("Area_list decode failed. Fix Area_list.npz decoding first.") if cluster_path.suffix.lower() == ".npz": with np.load(cluster_path, allow_pickle=True) as data: clustering = {k: data[k] for k in data.files} elif cluster_path.suffix.lower() == ".pkl": with open(cluster_path, "rb") as f: clustering = pickle.load(f) else: raise ValueError(f"Unsupported clustering file format: {cluster_path.suffix}") if not isinstance(clustering, dict): raise TypeError(f"Unexpected clustering payload type: {type(clustering)}") if "Cluster_Counter_Ordered" not in clustering: raise KeyError("Missing key 'Cluster_Counter_Ordered' in clustering data") if "Neuron_Counter_Ordered" not in clustering: raise KeyError("Missing key 'Neuron_Counter_Ordered' in clustering data") cluster_counter = np.asarray(clustering["Cluster_Counter_Ordered"]).reshape(-1) cluster_counter = np.asarray(np.round(cluster_counter), dtype=np.int32) neuron_counter = np.asarray(clustering["Neuron_Counter_Ordered"]).reshape(-1) neuron_counter = np.asarray(np.round(neuron_counter), dtype=np.int64) # Code_M's Data_Clustering.mat names this field 'Id_Ordered'; older Code_P # exports used lowercase 'Id_ordered'. Accept either so both load. id_key = "Id_Ordered" if "Id_Ordered" in clustering else ("Id_ordered" if "Id_ordered" in clustering else None) if id_key is not None: id_ordered = np.asarray(clustering[id_key]).reshape(-1) id_ordered = np.asarray(np.round(id_ordered), dtype=np.int64) else: if require_id_ordered: raise KeyError("Missing key 'Id_Ordered'/'Id_ordered' in clustering data.") id_ordered = neuron_counter.copy() # Materialize the unique cluster list once so downstream notebooks can use # it directly for ordered iteration and validation. unique_clusters = np.unique(cluster_counter) return { "entries_raw": entries_raw, "area_list": area_list, "enable_ccf_filter": enable_ccf_filter, "cluster_counter": cluster_counter, "neuron_counter": neuron_counter, "id_ordered": id_ordered, "unique_clusters": unique_clusters, } def load_canonical_movement_subspace_entries( psth_path: Path, mov_subspace_path: Path, subspace_key: str, ) -> Dict[str, Any]: """Load canonical PSTH entries and attach one movement subspace per entry. Parameters ---------- psth_path: NPZ file containing `psth_mat`. mov_subspace_path: NPZ file containing `mov_potent_null`. subspace_key: Subspace name to attach. Supported values are `null` and `potent`. Returns ------- dict[str, Any] Canonical entry bundle with one attached movement-subspace tensor per cleaned entry. """ subspace_key = str(subspace_key) if subspace_key not in {"null", "potent"}: raise ValueError("subspace_key must be 'null' or 'potent'") psth_path = Path(psth_path) mov_subspace_path = Path(mov_subspace_path) if not psth_path.exists(): raise FileNotFoundError(f"File not found: {psth_path}") if not mov_subspace_path.exists(): raise FileNotFoundError( f"Missing file: {mov_subspace_path}. Convert Movement_subspace.mat to Movement_subspace.npz first." ) with np.load(psth_path, allow_pickle=True) as data: if "psth_mat" not in data: raise KeyError("NPZ missing key 'psth_mat'") psth_arr = np.asarray(data["psth_mat"]).reshape(-1) with np.load(mov_subspace_path, allow_pickle=True) as data: if "mov_potent_null" not in data: raise KeyError("NPZ missing key 'mov_potent_null'") mov_arr = np.asarray(data["mov_potent_null"]).reshape(-1) psth_entries_raw = normalize_raw_entries(psth_arr) mov_entries_raw = normalize_raw_entries(mov_arr) if len(psth_entries_raw) != len(mov_entries_raw): raise RuntimeError( f"Entry count mismatch: psth={len(psth_entries_raw)} vs mov_potent_null={len(mov_entries_raw)}. " "MATLAB script assumes same probe indexing." ) id_ordered, neuron_counter = infer_placeholder_id_order(psth_entries_raw) canon = canonicalize_psth_entries_for_decoding( entries_raw=psth_entries_raw, id_ordered=id_ordered, neuron_counter=neuron_counter, ) entries_clean = canon["entries_clean"] # Attach the selected subspace tensor entry-by-entry so downstream decoding # code can stay agnostic to how the subspaces were originally stored. for i, entry in enumerate(entries_clean): mov_entry = mov_entries_raw[i] subspace_data = np.asarray(unwrap_scalar_obj(mov_entry.get(subspace_key, np.array([]))), dtype=np.float32) if subspace_data.size == 0: subspace_data = np.zeros((0, 0, 0), dtype=np.float32) if subspace_data.ndim != 3 and subspace_data.size > 0: raise ValueError(f"Entry {i}: {subspace_key} data must be 3D, got shape {subspace_data.shape}") entry[subspace_key] = subspace_data return { "psth_entries_raw": psth_entries_raw, "mov_entries_raw": mov_entries_raw, "entries_clean": entries_clean, "entries_by_area": canon["entries_by_area"], "window_centers": canon["window_centers"], "subspace_key": subspace_key, } def load_canonical_spontlick_entries( psth_path: Path, area_list_path: Path, lick_times_path: Path, ) -> Dict[str, Any]: """Load canonical PSTH entries and align spontaneous-lick metadata. Parameters ---------- psth_path: NPZ file containing `psth_mat`. area_list_path: NPZ file containing the area-list payload. lick_times_path: NPZ file containing `lick_times`. Returns ------- dict[str, Any] Canonical PSTH bundle augmented with cleaned spontaneous-lick metadata. """ psth_path = Path(psth_path) area_list_path = Path(area_list_path) lick_times_path = Path(lick_times_path) if not psth_path.exists(): raise FileNotFoundError(f"File not found: {psth_path}") if not lick_times_path.exists(): raise FileNotFoundError(f"File not found: {lick_times_path}") canon = load_canonical_psth_entries_and_area_list(psth_path, area_list_path) with np.load(lick_times_path, allow_pickle=True) as data: if "lick_times" not in data: raise KeyError("NPZ missing key 'lick_times'") lick_raw_arr = np.asarray(data["lick_times"]).reshape(-1) lick_entries_raw = normalize_raw_entries(lick_raw_arr) lick_clean: List[Dict[str, Any]] = [] for i, entry in enumerate(lick_entries_raw): session_id = to_session_id_str(entry.get("session_id", "")) lick_time = to_1d(entry.get("lick_time", np.array([], dtype=np.float32)), np.float32) video_ts = to_1d(entry.get("video_timestamp_continues", np.array([], dtype=np.float32)), np.float32) jaw_cont = to_1d(entry.get("jaw_movement_continues", np.array([], dtype=np.float32)), np.float32) lick_mask = to_1d(entry.get("lick_mask", np.array([], dtype=bool))).astype(bool) n = min(video_ts.size, jaw_cont.size, lick_mask.size) if n > 0: video_ts = video_ts[:n] jaw_cont = jaw_cont[:n] lick_mask = lick_mask[:n] lick_clean.append( { "entry_idx": i, "session_id": session_id, "lick_time": lick_time, "video_timestamp_continues": video_ts, "jaw_movement_continues": jaw_cont, "lick_mask": lick_mask, } ) lick_by_idx = {entry["entry_idx"]: entry for entry in lick_clean} lick_by_session: Dict[str, List[Dict[str, Any]]] = {} for entry in lick_clean: lick_by_session.setdefault(entry["session_id"], []).append(entry) # Prefer entry-index alignment first, then fall back to session-level # matching when the spontaneous-lick export uses a different local order. for entry in canon["entries_clean"]: lick_entry = lick_by_idx.get(entry["entry_idx"]) if (lick_entry is None) or (lick_entry["session_id"] != entry["session_id"]): candidates = lick_by_session.get(entry["session_id"], []) lick_entry = candidates[0] if candidates else None entry["lick_entry"] = lick_entry return { "psth_entries_raw": canon["entries_raw"], "lick_entries_raw": lick_entries_raw, "entries_clean": canon["entries_clean"], "entries_by_area": canon["entries_by_area"], "window_centers": canon["window_centers"], "area_list": canon["area_list"], "lick_clean": lick_clean, } def canonicalize_psth_entries_for_decoding( entries_raw: List[Dict[str, Any]], id_ordered: np.ndarray, neuron_counter: np.ndarray, ) -> Dict[str, Any]: """Canonicalize raw PSTH entries for decoding workflows. Parameters ---------- entries_raw: Raw PSTH entry payloads loaded from disk. id_ordered: Global unit-order vector used to validate canonical IDs. neuron_counter: Per-unit ordering vector carried alongside `id_ordered`. Returns ------- dict[str, Any] Canonical cleaned entries, entries grouped by area, and canonical window centers. """ entries_by_area: Dict[str, List[Dict[str, Any]]] = {} entries_clean: List[Dict[str, Any]] = [] # Rebuild each entry with one stable set of field names and shapes so the # decoding stack does not have to reason about multiple raw payload layouts. for i, entry in enumerate(entries_raw): spike_counts = np.asarray(unwrap_scalar_obj(entry["spike_counts"]), dtype=np.float32) if spike_counts.ndim != 3: raise ValueError(f"Entry {i}: spike_counts must be 3D, got {spike_counts.shape}") _, n_trials, n_units = spike_counts.shape trial = _normalize_vector_length( _entry_get_any(entry, ("trial_type", "trial"), np.zeros(n_trials, dtype=np.float32)), expected_size=n_trials, dtype=np.float32, fill_value=0, ) lick = _normalize_vector_length( _entry_get_any(entry, ("lick_flag", "lick"), np.zeros(n_trials, dtype=bool)), expected_size=n_trials, dtype=bool, fill_value=False, ) stim = _normalize_vector_length( _entry_get_any(entry, ("whisker_stim", "stim"), np.zeros(n_trials, dtype=bool)), expected_size=n_trials, dtype=bool, fill_value=False, ) early_lick = _normalize_vector_length( _entry_get_any(entry, ("early_lick",), np.zeros(n_trials, dtype=bool)), expected_size=n_trials, dtype=bool, fill_value=False, ) quiet_whisker = _normalize_vector_length( _entry_get_any(entry, ("quiet_trial_whisker_speed", "quiet_whisker"), np.ones(n_trials, dtype=bool)), expected_size=n_trials, dtype=bool, fill_value=True, ) quiet_jaw = _normalize_vector_length( _entry_get_any(entry, ("quiet_trial_jaw_movement", "quiet_jaw"), np.ones(n_trials, dtype=bool)), expected_size=n_trials, dtype=bool, fill_value=True, ) lick_time = _normalize_vector_length( _entry_get_any(entry, ("lick_time",), np.zeros(n_trials, dtype=np.float32)), expected_size=n_trials, dtype=np.float32, fill_value=0.0, ) start_time = _normalize_vector_length( _entry_get_any(entry, ("start_time",), np.zeros(n_trials, dtype=np.float32)), expected_size=n_trials, dtype=np.float32, fill_value=0.0, ) if trial.size != n_trials or lick.size != n_trials: raise ValueError(f"Entry {i}: trial/lick size mismatch with n_trials={n_trials}") unit_rs = _normalize_vector_length( _entry_get_any(entry, ("unit_rsUnits", "unit_rs"), np.ones(n_units, dtype=bool)), expected_size=n_units, dtype=bool, fill_value=True, ) unit_fs = _normalize_vector_length( _entry_get_any(entry, ("unit_fsUnits", "unit_fs"), np.ones(n_units, dtype=bool)), expected_size=n_units, dtype=bool, fill_value=True, ) unit_ccf_location = _normalize_vector_length( to_1d_str(_entry_get_any(entry, ("unit_ccf_location",), np.array([""] * n_units, dtype=object))), expected_size=n_units, dtype=object, fill_value="", ) unit_spike_times_raw = unwrap_scalar_obj(_entry_get_any(entry, ("unit_spike_times",), np.array([], dtype=object))) unit_spike_times = np.asarray(unit_spike_times_raw, dtype=object).reshape(-1) # Spike times are absolute session timestamps (up to ~10^4 s). Code_M keeps # them in double precision; float32 (~0.5 ms near 10^4 s) would shift # spikes across bin edges whenever they are re-binned relative to an event # (e.g. spontaneous-lick coding direction). Keep float64. if unit_spike_times.size != n_units: tmp = np.empty(n_units, dtype=object) for k in range(n_units): tmp[k] = np.array([], dtype=np.float64) for k in range(min(n_units, unit_spike_times.size)): tmp[k] = np.asarray(unwrap_scalar_obj(unit_spike_times[k]), dtype=np.float64).reshape(-1) unit_spike_times = tmp else: for k in range(n_units): unit_spike_times[k] = np.asarray(unwrap_scalar_obj(unit_spike_times[k]), dtype=np.float64).reshape(-1) globalcluster = _normalize_vector_length( _entry_get_any(entry, ("GlobalclusterID",), np.arange(1, n_units + 1)), expected_size=n_units, dtype=np.int64, fill_value=0, ) if not np.any(globalcluster): globalcluster = np.arange(1, n_units + 1, dtype=np.int64) trial_timestamps = _normalize_vector_length( _entry_get_any(entry, ("trial_timestamps",), np.arange(spike_counts.shape[0], dtype=np.float32)), expected_size=spike_counts.shape[0], dtype=np.float32, fill_value=np.nan, ) if not np.any(np.isfinite(trial_timestamps)): trial_timestamps = np.arange(spike_counts.shape[0], dtype=np.float32) else: trial_timestamps = sanitize_time_axis(trial_timestamps) record = { "canonical_api_version": FUNCTION_API_VERSION, "entry_idx": i, "session_id": to_session_id_str(entry.get("session_id", f"entry_{i}")), "probe_location": to_session_id_str(entry.get("probe_location", "")), "spike_counts": spike_counts, "trial": trial, "trial_type": trial, "lick": lick, "lick_flag": lick, "stim": stim, "whisker_stim": stim, "early_lick": early_lick, "quiet_whisker": quiet_whisker, "quiet_trial_whisker_speed": quiet_whisker, "quiet_jaw": quiet_jaw, "quiet_trial_jaw_movement": quiet_jaw, "lick_time": lick_time, "start_time": start_time, "unit_rs": unit_rs, "unit_rsUnits": unit_rs, "unit_fs": unit_fs, "unit_fsUnits": unit_fs, "unit_ccf_location": unit_ccf_location, "unit_spike_times": unit_spike_times, "GlobalclusterID": globalcluster, "trial_timestamps": trial_timestamps, "n_trials": n_trials, "n_units": n_units, } entries_clean.append(record) entries_by_area.setdefault(record["probe_location"], []).append(record) all_global_ids = np.concatenate( [entry["GlobalclusterID"].reshape(-1) for entry in entries_clean if entry["GlobalclusterID"].size > 0] ).astype(np.int64, copy=False) if all_global_ids.size > 0: gid_max = int(np.max(all_global_ids)) id_max = int(np.max(id_ordered)) if id_ordered.size else 0 if np.array_equal(id_ordered, neuron_counter) and (gid_max > id_max): raise RuntimeError( f"Id_ordered appears to be Neuron_Counter_Ordered (1..{id_max}) instead of GlobalclusterID/Unit_ids " f"(global max={gid_max}). Re-run Data_Clustering export so Id_ordered uses global Unit_ids." ) sample_area_entries = next(iter(entries_by_area.values())) if entries_by_area else [] if not sample_area_entries: raise RuntimeError("No probe entries found in PSTH data") window_centers = sanitize_time_axis(sample_area_entries[0]["trial_timestamps"]) return { "entries_by_area": entries_by_area, "entries_clean": entries_clean, "window_centers": window_centers, }