"""Public analysis module. Conceptually, this module evaluates decoding performance after removing either identified clusters or matched random neuron subsets from the canonical entry set. It exists as a separate unit so dropout bookkeeping, unit masking, and bin-wise decoding stay consistent across the structured dropout analyses. It connects clustered neural entries to the per-area accuracy bundles used by the dropout notebooks and summary figures. """ 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_ccf_mask, get_celltype_mask, get_completion_mask, get_quiet_mask, ) def run_cluster_dropout_decoding( entries_by_area: Mapping[str, List[Dict[str, Any]]], unique_clusters: np.ndarray, cluster_counter: np.ndarray, neuron_counter: np.ndarray, params_drop: Mapping[str, Any], id_ordered: Optional[np.ndarray] = None, area_list: Optional[Mapping[str, Set[str]]] = None, enable_ccf_filter: bool = True, rng_seed: int = 0, ) -> Dict[str, Any]: """Run decoding after removing one identified cluster at a time. Parameters ---------- entries_by_area: Canonical entries grouped by area. unique_clusters: Cluster ids evaluated independently. cluster_counter, neuron_counter: Parallel arrays mapping each neuron id to its cluster id. id_ordered: Global unit ids (`Id_Ordered`) aligned with `cluster_counter`. These are the ids that match `GlobalclusterID` on the PSTH entries. params_drop: Decoding configuration mapping containing region list, trial filters, window centers, and SVM settings. area_list: Optional area-to-CCF mapping used by the unit filter. enable_ccf_filter: Whether CCF-based unit filtering is active. rng_seed: Seed for the shared RNG driving balancing and shuffle controls. Returns ------- dict[str, Any] Cluster-keyed accuracy bundles with real and shuffled decoding scores. """ rng = np.random.default_rng(rng_seed) acc_dropout: Dict[str, Any] = {} for i_cluster in np.asarray(unique_clusters).reshape(-1): cluster_name = f"cluster{int(i_cluster)}" id_2drop = np.asarray(id_ordered)[np.asarray(cluster_counter) == i_cluster] accuracy: Dict[str, Any] = {"sessionaddress": {}} accuracy_shuffeled: Dict[str, Any] = {"sessionaddress": {}} for currentarea in params_drop["regionlist"]: probe_entries = list(entries_by_area.get(currentarea, [])) area_acc_list: List[np.ndarray] = [] area_sh_list: List[np.ndarray] = [] area_sessions: List[str] = [] for e in probe_entries: trial = e["trial"] lick = e["lick"] completed_trials_ind = get_completion_mask(e, params_drop["completion_state"]) qind = get_quiet_mask(e, params_drop["quietstate"]) class1 = (trial == 1) & (lick == 1) class2 = (trial == 3) & (lick == 0) class1 = class1 & qind & completed_trials_ind class2 = class2 & qind & completed_trials_ind celltype_ind = get_celltype_mask(e, params_drop["celltype"]) ccf_ind = get_ccf_mask( e, currentarea, area_list=area_list, enable_ccf_filter=enable_ccf_filter, ) # Remove all units whose global ids belong to the current # cluster before the usual area and cell-type filters are applied. curr_id_not2drop = ~np.isin(e["GlobalclusterID"], id_2drop) curr_cell_ind = celltype_ind & ccf_ind & curr_id_not2drop if np.sum(curr_cell_ind) < 5: continue currsig = e["spike_counts"][:, :, curr_cell_ind] val_class1 = currsig[:, class1, :] val_class2 = currsig[:, class2, :] n_bins = currsig.shape[0] sess_acc = np.full((n_bins,), np.nan, dtype=np.float32) sess_sh = np.full((n_bins,), np.nan, dtype=np.float32) n1 = int(np.sum(class1)) n2 = int(np.sum(class2)) for iBin in range(n_bins): X1 = np.squeeze(val_class1[iBin, :, :]) X2 = np.squeeze(val_class2[iBin, :, :]) X1 = ensure_trials_by_cells(X1, n1) X2 = ensure_trials_by_cells(X2, n2) acc_val, sh_val = decode_one_bin_svm_matlab_compat( X1, X2, mintrial=params_drop["mintrial"], balance_method=params_drop["balance_method"], zscoring=bool(params_drop["zscoring"]), run_seed=0, rng_obj=rng, ) sess_acc[iBin] = acc_val sess_sh[iBin] = sh_val area_acc_list.append(sess_acc) area_sh_list.append(sess_sh) area_sessions.append(e["session_id"]) if area_acc_list: accuracy[currentarea] = np.stack(area_acc_list, axis=0) accuracy_shuffeled[currentarea] = np.stack(area_sh_list, axis=0) else: n_bins = int(np.asarray(params_drop["windowCenters"]).size) accuracy[currentarea] = np.zeros((0, n_bins), dtype=np.float32) accuracy_shuffeled[currentarea] = np.zeros((0, n_bins), dtype=np.float32) accuracy["sessionaddress"][currentarea] = np.array(area_sessions, dtype=object) accuracy_shuffeled["sessionaddress"][currentarea] = np.array(area_sessions, dtype=object) acc_dropout[cluster_name] = { "Accuracy": accuracy, "Accuracy_shuffeled": accuracy_shuffeled, "windowCenters": params_drop["windowCenters"], } return acc_dropout def run_random_dropout_decoding( entries_by_area: Mapping[str, List[Dict[str, Any]]], unique_clusters: np.ndarray, cluster_counter: np.ndarray, id_ordered: np.ndarray, params_rnd: Mapping[str, Any], area_list: Optional[Mapping[str, Set[str]]] = None, enable_ccf_filter: bool = True, rng_seed: int = 0, ) -> Dict[str, Any]: """Run decoding after removing matched random neuron subsets. Parameters ---------- entries_by_area: Canonical entries grouped by area. unique_clusters: Cluster ids whose population sizes define the matched random dropouts. cluster_counter: Cluster label per globally ordered neuron id. id_ordered: Global neuron ids aligned to `cluster_counter`. params_rnd: Decoding configuration mapping containing region list, trial filters, window centers, and SVM settings. area_list: Optional area-to-CCF mapping used by the unit filter. enable_ccf_filter: Whether CCF-based unit filtering is active. rng_seed: Seed for the shared RNG driving balancing and shuffle controls. Returns ------- dict[str, Any] Cluster-keyed accuracy bundles for the matched random-dropout control. """ rng = np.random.default_rng(rng_seed) acc_random: Dict[str, Any] = {} for i_cluster in np.asarray(unique_clusters).reshape(-1): cluster_name = f"cluster{int(i_cluster)}" id_2drop = np.asarray(id_ordered)[np.asarray(cluster_counter) == i_cluster] accuracy_rnd_selection: Dict[str, Any] = {"sessionaddress": {}} accuracy_rnd_selection_shuffeled: Dict[str, Any] = {"sessionaddress": {}} for currentarea in params_rnd["regionlist"]: probe_entries = list(entries_by_area.get(currentarea, [])) area_acc_list: List[np.ndarray] = [] area_sh_list: List[np.ndarray] = [] area_sessions: List[str] = [] for e in probe_entries: trial = e["trial"] lick = e["lick"] completed_trials_ind = get_completion_mask(e, params_rnd["completion_state"]) qind = get_quiet_mask(e, params_rnd["quietstate"]) class1 = (trial == 1) & (lick == 1) class2 = (trial == 3) & (lick == 0) class1 = class1 & qind & completed_trials_ind class2 = class2 & qind & completed_trials_ind celltype_ind = get_celltype_mask(e, params_rnd["celltype"]) ccf_ind = get_ccf_mask( e, currentarea, area_list=area_list, enable_ccf_filter=enable_ccf_filter, ) curr_cell_ind = celltype_ind & ccf_ind curr_global_id = e["GlobalclusterID"] curr_id_2drop = np.where(np.isin(curr_global_id, id_2drop))[0] if not np.all(np.isin(curr_id_2drop, np.where(curr_cell_ind)[0])): bad_idx = curr_id_2drop[~np.isin(curr_id_2drop, np.where(curr_cell_ind)[0])] raise ValueError( f"Stop here: something is wrong with cluster dropout | cluster={cluster_name} " f"area={currentarea} session={e['session_id']} bad_idx_count={bad_idx.size} " f"first_bad={bad_idx[:10].tolist()} id_ordered_range=[{int(np.min(id_ordered))},{int(np.max(id_ordered))}]" ) curr_id_not2drop = np.ones(e["n_units"], dtype=bool) curr_id_not2drop[curr_id_2drop] = False num_vectors = 5 num_zeros = int(np.sum(~curr_id_not2drop)) available_positions = np.where(curr_id_not2drop & curr_cell_ind)[0] n_bins = e["spike_counts"].shape[0] sess_acc = np.full((n_bins, num_vectors), np.nan, dtype=np.float32) sess_sh = np.full((n_bins, num_vectors), np.nan, dtype=np.float32) n1 = int(np.sum(class1)) n2 = int(np.sum(class2)) # Keep the cumulative masking behavior of the original random # control so each draw removes additional cells from the same # working subset within one probe. curr_cell_ind_rnd = curr_cell_ind.copy() for rnd_selection in range(num_vectors): if num_zeros > 0: if available_positions.size < num_zeros: continue random_zero_positions = rng.choice(available_positions, size=num_zeros, replace=False) else: random_zero_positions = np.array([], dtype=np.int64) temp_vec = np.ones(e["n_units"], dtype=bool) temp_vec[random_zero_positions] = False curr_cell_ind_rnd = curr_cell_ind_rnd & temp_vec if np.sum(curr_cell_ind_rnd) < 5: continue currsig = e["spike_counts"][:, :, curr_cell_ind_rnd] val_class1 = currsig[:, class1, :] val_class2 = currsig[:, class2, :] for iBin in range(n_bins): X1 = np.squeeze(val_class1[iBin, :, :]) X2 = np.squeeze(val_class2[iBin, :, :]) X1 = ensure_trials_by_cells(X1, n1) X2 = ensure_trials_by_cells(X2, n2) acc_val, sh_val = decode_one_bin_svm_matlab_compat( X1, X2, mintrial=params_rnd["mintrial"], balance_method=params_rnd["balance_method"], zscoring=bool(params_rnd["zscoring"]), run_seed=0, rng_obj=rng, ) sess_acc[iBin, rnd_selection] = acc_val sess_sh[iBin, rnd_selection] = sh_val area_acc_list.append(sess_acc) area_sh_list.append(sess_sh) area_sessions.append(e["session_id"]) if area_acc_list: accuracy_rnd_selection[currentarea] = np.stack(area_acc_list, axis=0) accuracy_rnd_selection_shuffeled[currentarea] = np.stack(area_sh_list, axis=0) else: n_bins = int(np.asarray(params_rnd["windowCenters"]).size) accuracy_rnd_selection[currentarea] = np.zeros((0, n_bins, 5), dtype=np.float32) accuracy_rnd_selection_shuffeled[currentarea] = np.zeros((0, n_bins, 5), dtype=np.float32) accuracy_rnd_selection["sessionaddress"][currentarea] = np.array(area_sessions, dtype=object) accuracy_rnd_selection_shuffeled["sessionaddress"][currentarea] = np.array(area_sessions, dtype=object) acc_random[cluster_name] = { "Accuracy_rnd_selection": accuracy_rnd_selection, "Accuracy_rnd_selection_shuffeled": accuracy_rnd_selection_shuffeled, "windowCenters": params_rnd["windowCenters"], } return acc_random