"""Public helper module. Conceptually, this module groups the lightweight coercion helpers used by structured decoding figure notebooks. It exists as a separate unit so notebooks can stay linear while the shape normalization and session filtering logic lives in reusable Python helpers. It connects saved decoding accuracy payloads to time-by-session matrices that can be summarized and plotted directly. """ from __future__ import annotations from pathlib import Path from typing import Any import numpy as np from .core import to_session_id_str, unwrap_scalar_obj from .figure_data import to_dict_like from .signal_utils import nearest_bin __all__ = [ "align_decoding_matrices_by_session", "extract_decoding_signal_matrix", "keep_nonzero_paired_columns", "load_decoding_npz", "nonempty_indices", "session_ids_for_area", "window_index_bounds", ] def load_decoding_npz(path: str | Path) -> tuple[dict[str, Any], np.ndarray]: """Load one saved decoding payload and its time-bin vector. The processed decoding `.npz` files store one top-level `Accuracy` payload as a nested scalar object. This helper normalizes that payload into a plain dict and returns the associated `windowCenters` vector, falling back to the MATLAB-style default grid when that key is absent. """ with np.load(Path(path), allow_pickle=True) as data: if "Accuracy" not in data: raise KeyError(f"{Path(path).name} missing key 'Accuracy'") accuracy = to_dict_like(np.asarray(data["Accuracy"]).reshape(-1)[0]) if "windowCenters" in data: window_centers = np.asarray(data["windowCenters"], dtype=float).reshape(-1) else: window_centers = np.arange(-0.99, 2.01, 0.01, dtype=float) return accuracy, window_centers def _is_empty_like(value: Any) -> bool: """Return whether a nested scalar-like payload carries no usable content.""" unpacked = unwrap_scalar_obj(value) if unpacked is None: return True if isinstance(unpacked, (str, bytes)): return len(unpacked) == 0 if isinstance(unpacked, np.ndarray): if unpacked.size == 0: return True if unpacked.dtype == object and unpacked.shape == (): return _is_empty_like(unpacked.item()) return False if isinstance(unpacked, (list, tuple, dict, set)): return len(unpacked) == 0 return False def nonempty_indices(session_address_area: Any) -> np.ndarray: """Return indices of non-empty session-address entries for one area.""" unpacked = unwrap_scalar_obj(session_address_area) if isinstance(unpacked, np.ndarray): items = unpacked.reshape(-1).tolist() elif isinstance(unpacked, (list, tuple)): items = list(unpacked) else: if _is_empty_like(unpacked): return np.asarray([], dtype=int) return np.asarray([0], dtype=int) keep = [index for index, item in enumerate(items) if not _is_empty_like(item)] return np.asarray(keep, dtype=int) def session_ids_for_area(acc_struct: Any, area: str) -> list[str]: """Return one session-id string per stored session-address entry.""" acc = to_dict_like(acc_struct) if "sessionaddress" not in acc: return [] sessionaddress = to_dict_like(acc["sessionaddress"]) if area not in sessionaddress: return [] arr = np.asarray(unwrap_scalar_obj(sessionaddress[area]), dtype=object).reshape(-1) return [to_session_id_str(value) for value in arr] def align_decoding_matrices_by_session( matrix_a: Any, session_ids_a: list[str], matrix_b: Any, session_ids_b: list[str], ) -> tuple[np.ndarray, np.ndarray, bool]: """Align two time-by-session matrices by non-empty session ids. Returns the input matrices unchanged with `False` when the provided id vectors do not match the matrix widths. Otherwise, returns the aligned matrices and `True`, even when no shared non-empty ids remain. """ mat_a = np.asarray(matrix_a, dtype=float) mat_b = np.asarray(matrix_b, dtype=float) if mat_a.ndim != 2 or mat_b.ndim != 2: raise ValueError("matrix_a and matrix_b must both be two-dimensional") if len(session_ids_a) != mat_a.shape[1] or len(session_ids_b) != mat_b.shape[1]: return mat_a, mat_b, False index_b_by_id: dict[str, int] = {} for index_b, session_id in enumerate(session_ids_b): if session_id and session_id not in index_b_by_id: index_b_by_id[session_id] = index_b keep_a: list[int] = [] keep_b: list[int] = [] for index_a, session_id in enumerate(session_ids_a): if not session_id or session_id not in index_b_by_id: continue keep_a.append(index_a) keep_b.append(index_b_by_id[session_id]) if len(keep_a) == 0: return mat_a[:, :0], mat_b[:, :0], True return mat_a[:, keep_a], mat_b[:, keep_b], True def keep_nonzero_paired_columns(matrix_a: Any, matrix_b: Any) -> tuple[np.ndarray, np.ndarray]: """Drop paired session columns that are all-zero in either matrix.""" mat_a = np.asarray(matrix_a, dtype=float) mat_b = np.asarray(matrix_b, dtype=float) if mat_a.ndim != 2 or mat_b.ndim != 2: raise ValueError("matrix_a and matrix_b must both be two-dimensional") n_cols = min(mat_a.shape[1], mat_b.shape[1]) mat_a = mat_a[:, :n_cols] mat_b = mat_b[:, :n_cols] if n_cols == 0: return mat_a, mat_b all_zero_a = np.all(mat_a == 0, axis=0) all_zero_b = np.all(mat_b == 0, axis=0) keep = ~(all_zero_a | all_zero_b) return mat_a[:, keep], mat_b[:, keep] def window_index_bounds(window_centers: Any, bounds: tuple[float, float]) -> tuple[int, int]: """Return inclusive nearest-bin bounds for a requested time window.""" centers = np.asarray(window_centers, dtype=float).reshape(-1) if centers.size == 0: raise ValueError("window_centers cannot be empty") first = int(nearest_bin(centers, float(bounds[0]))) last = int(nearest_bin(centers, float(bounds[1]))) if last < first: first, last = last, first return first, last def extract_decoding_signal_matrix( current_nb_mat: Any, area: str, window_centers: Any, keep_indices: Any = None, ) -> np.ndarray: """Return a time-by-session decoding matrix for one area. The saved decoding payloads can arrive as 1D, 2D, or 3D arrays depending on whether they already contain per-session means or still carry a unit axis. This helper normalizes those variants into one common `(n_time_bins, n_sessions)` representation and then filters to the subset of sessions with non-empty addresses. """ nb_mat = to_dict_like(current_nb_mat) centers = np.asarray(window_centers, dtype=float).reshape(-1) n_time = int(centers.size) if n_time == 0: raise ValueError("window_centers cannot be empty") if area not in nb_mat: return np.zeros((n_time, 0), dtype=float) array = np.asarray(unwrap_scalar_obj(nb_mat[area]), dtype=float) if array.size == 0: return np.zeros((n_time, 0), dtype=float) if array.ndim == 3: finite_counts = np.sum(np.isfinite(array), axis=2) sums = np.nansum(array, axis=2) signal_2d = np.full(sums.shape, np.nan, dtype=float) valid = finite_counts > 0 signal_2d[valid] = sums[valid] / finite_counts[valid] signal = signal_2d.T elif array.ndim == 2: signal = array.T elif array.ndim == 1: signal = array.reshape(-1, 1) else: signal = array.reshape(array.shape[0], -1) if signal.shape[0] != n_time and signal.shape[1] == n_time: signal = signal.T if signal.shape[0] != n_time: raise ValueError( f"Decoded signal for area {area!r} has shape {signal.shape}, " f"which does not match {n_time} window centers" ) if keep_indices is None: return signal keep = np.asarray(keep_indices, dtype=int).reshape(-1) keep = keep[(keep >= 0) & (keep < signal.shape[1])] if keep.size == 0: return signal[:, :0] return signal[:, keep]