"""Public analysis module. Conceptually, this module evaluates decoding performance inside the learned movement subspaces used by the analysis notebooks. It exists as a separate unit so projected-subspace payloads, trial selection, and bin-wise decoding stay synchronized across potent and null analyses. It connects subspace-enriched entries to notebook-ready per-area accuracy matrices. """ from __future__ import annotations from typing import Any, Dict, List, Mapping, Optional, Set, Tuple import numpy as np from .loading import _entry_get_any, _infer_entry_n_trials, _infer_entry_n_units, _normalize_vector_length from .math_utils import ( downsample_balance, ensure_trials_by_cells, gram_schmidt_columns, holdout_split, nearest_bin, normalize_vec, project_trials, psth_simple_counts_single_event, zscore_cols, ) from .decoding import decode_one_bin_svm_matlab_compat, get_celltype_mask, get_completion_mask, get_quiet_mask def compute_condition_projection_trace(subspace_tensor: Any, trial_mask: Any) -> np.ndarray: """Return one movement-subspace magnitude trace for a selected trial subset. Parameters ---------- subspace_tensor: Three-dimensional projected activity with shape `(n_bins, n_trials, n_dims)`. trial_mask: Boolean mask selecting the trials that contribute to the condition average. Returns ------- np.ndarray One-dimensional trace with shape `(n_bins,)`. Returns all-`NaN` when the requested condition has no surviving trials. """ subspace_tensor = np.asarray(subspace_tensor, dtype=np.float32) trial_mask = np.asarray(trial_mask, dtype=bool).reshape(-1) if subspace_tensor.ndim != 3: raise ValueError(f"subspace_tensor must be 3D, got {subspace_tensor.shape}") n_bins, n_trials, _ = subspace_tensor.shape if trial_mask.size != n_trials: raise ValueError("trial_mask size mismatch with subspace tensor") if not np.any(trial_mask): return np.full((n_bins,), np.nan, dtype=np.float32) mean_over_trials = np.nanmean(subspace_tensor[:, trial_mask, :], axis=1) return np.nansum(mean_over_trials**2, axis=1).astype(np.float32) def collect_area_condition_projection_matrices( entries_potent: list[dict[str, Any]], entries_null: list[dict[str, Any]], *, current_area: str, trial_types: list[int] | tuple[int, ...], lick_states: list[bool] | tuple[bool, ...], completion_state: str = "completed_trials", n_bins: int | None = None, ) -> dict[str, list[np.ndarray]]: """Collect one potent/null session matrix per condition for one area. Parameters ---------- entries_potent, entries_null: Canonical PSTH entries enriched with the corresponding potent or null movement-subspace tensor. current_area: Probe location retained for the current panel. trial_types, lick_states: Parallel condition definitions. Each pair selects one trial subset. completion_state: Trial-completion selector forwarded to `get_completion_mask`. n_bins: Optional fallback number of bins used when no session contributes to a condition. Returns ------- dict[str, list[np.ndarray]] Two lists of condition matrices keyed by `potent` and `null`, with one `(n_bins, n_sessions)` matrix per condition. """ if len(entries_potent) != len(entries_null): raise ValueError("entries_potent and entries_null must have the same length") if len(trial_types) != len(lick_states): raise ValueError("trial_types and lick_states must have the same length") potent_by_cond: list[list[np.ndarray]] = [[] for _ in trial_types] null_by_cond: list[list[np.ndarray]] = [[] for _ in trial_types] inferred_n_bins = int(n_bins or 0) for entry_potent, entry_null in zip(entries_potent, entries_null): if str(entry_potent.get("probe_location", "")) != str(current_area): continue current_potent = np.asarray(entry_potent.get("potent", np.zeros((0, 0, 0), dtype=np.float32)), dtype=np.float32) current_null = np.asarray(entry_null.get("null", np.zeros((0, 0, 0), dtype=np.float32)), dtype=np.float32) if current_potent.size == 0 or current_null.size == 0: continue inferred_n_bins = max(inferred_n_bins, int(current_potent.shape[0])) trial = np.asarray(entry_potent["trial"]).reshape(-1) lick = np.asarray(entry_potent["lick"], dtype=bool).reshape(-1) completion = get_completion_mask(entry_potent, str(completion_state)) # Build one condition-specific session trace so the area matrices stay # session-aligned across potent and null projections. for condition_index, (trial_type, lick_state) in enumerate(zip(trial_types, lick_states)): curr_trial_ind = completion & (trial == float(trial_type)) & (lick == bool(lick_state)) potent_by_cond[condition_index].append(compute_condition_projection_trace(current_potent, curr_trial_ind)) null_by_cond[condition_index].append(compute_condition_projection_trace(current_null, curr_trial_ind)) potent_mats = [np.column_stack(v) if len(v) else np.zeros((inferred_n_bins, 0), dtype=np.float32) for v in potent_by_cond] null_mats = [np.column_stack(v) if len(v) else np.zeros((inferred_n_bins, 0), dtype=np.float32) for v in null_by_cond] return {"potent": potent_mats, "null": null_mats} def run_movement_subspace_decoding_pipeline( entries_by_area: Mapping[str, List[Dict[str, Any]]], params: Mapping[str, Any], subspace_key: str, rng_seed: int = 0, ) -> Dict[str, Any]: """Run the shared movement-subspace decoding pipeline. Parameters ---------- entries_by_area: Probe entries grouped by area. params: Decoding configuration mapping containing region list, window centers, trial filters, and SVM settings. subspace_key: Name of the projected subspace stored in each entry. Supported values are `"null"` and `"potent"`. rng_seed: Seed for the shared RNG driving balancing and shuffle controls. Returns ------- dict[str, Any] Accuracy bundle with real and shuffled scores plus the shared window centers. """ subspace_key = str(subspace_key) if subspace_key not in {"null", "potent"}: raise ValueError("subspace_key must be 'null' or 'potent'") rng = np.random.default_rng(rng_seed) window_centers = np.asarray(params["windowCenters"], dtype=np.float32) accuracy: Dict[str, Any] = {"sessionaddress": {}} accuracy_shuffeled: Dict[str, Any] = {"sessionaddress": {}} for current_area in params["regionlist"]: probe_entries = list(entries_by_area.get(current_area, [])) n_bins = int(window_centers.size) # Use the first non-empty projected tensor to size the per-area output # exactly like the stored subspace payload. for e in probe_entries: subspace_data = e.get(subspace_key, np.zeros((0, 0, 0), dtype=np.float32)) if isinstance(subspace_data, np.ndarray) and subspace_data.ndim == 3 and subspace_data.shape[0] > 0: n_bins = int(subspace_data.shape[0]) break area_acc = np.full((len(probe_entries), n_bins), np.nan, dtype=np.float32) area_sh = np.full((len(probe_entries), n_bins), np.nan, dtype=np.float32) area_sessions = np.full((len(probe_entries),), None, dtype=object) for row_idx, e in enumerate(probe_entries): trial = np.asarray(e["trial"]).reshape(-1) lick = np.asarray(e["lick"]).reshape(-1) completed_trials_ind = get_completion_mask(e, str(params["completion_state"])) # Preserved behavior: these masks are computed but not applied. _qind_unused = get_quiet_mask(e, str(params["quietstate"])) _celltype_unused = get_celltype_mask(e, str(params["celltype"])) class1 = (trial == 1) & (lick == 1) & completed_trials_ind class2 = (trial == 3) & (lick == 0) & completed_trials_ind currsig = e.get(subspace_key, np.zeros((0, 0, 0), dtype=np.float32)) if currsig.size == 0 or currsig.ndim != 3: continue val_class1 = currsig[:, class1, :] val_class2 = currsig[:, class2, :] n1 = int(np.sum(class1)) n2 = int(np.sum(class2)) # Decode each projected time bin independently so the output keeps # one accuracy trace per probe entry. for i_bin in range(min(n_bins, currsig.shape[0])): X1 = np.squeeze(val_class1[i_bin, :, :]) X2 = np.squeeze(val_class2[i_bin, :, :]) X1 = ensure_trials_by_cells(X1, n1) X2 = ensure_trials_by_cells(X2, n2) try: acc_val, sh_val = decode_one_bin_svm_matlab_compat( X1, X2, mintrial=int(params["mintrial"]), balance_method=str(params["balance_method"]), zscoring=bool(params["zscoring"]), rng_obj=rng, ) except ValueError: acc_val, sh_val = np.nan, np.nan area_acc[row_idx, i_bin] = np.float32(acc_val) area_sh[row_idx, i_bin] = np.float32(sh_val) area_sessions[row_idx] = e["session_id"] accuracy[current_area] = area_acc accuracy_shuffeled[current_area] = area_sh accuracy["sessionaddress"][current_area] = area_sessions accuracy_shuffeled["sessionaddress"][current_area] = area_sessions return { "Accuracy": accuracy, "Accuracy_shuffeled": accuracy_shuffeled, "windowCenters": window_centers, }