Analysis¶

def agg_by_neuron(
    y,
    neurons: pd.DataFrame,
    agg: Literal["mean", "sum", "std"] = "mean",
    neuron_type_column: str = "cell_type",
    **kwargs,
) -> dict:
    """Aggregate data by neuron type."""
    agg_func = getattr(np, agg) if isinstance(y, np.ndarray) else getattr(torch, agg)
    ret = {}
    for neuron_type, group in neurons.groupby(
        neuron_type_column, dropna=True, **kwargs
    ):
        ret[neuron_type] = agg_func(y[..., group.simple_id.to_numpy()], -1)
    return ret

def agg_by_neuropil(
    y,
    neurons: pd.DataFrame | None = None,
    connections: pd.DataFrame | None = None,
    mode: Literal["top_innervated", "all_innervated"] = "all_innervated",
    agg: Literal["mean", "sum", "std"] = "mean",
    use_polars: bool = False,
):
    """Aggregate activations by neuropil under a validated aggregation mode."""
    agg_func = getattr(np, agg) if isinstance(y, np.ndarray) else getattr(torch, agg)
    if use_polars:
        try:
            import polars as pl
        except ImportError:
            use_polars = False

    if mode == "top_innervated":
        assert neurons is not None, "neurons must be provided for top_innervated mode"
        tmp = neurons[["group", "simple_id"]].copy()
        tmp = tmp[tmp["simple_id"] < y.shape[-1]]
        pre_ret: dict = {}
        post_ret: dict = {}

        if use_polars:
            tmp["pre"] = tmp["group"].str.split(".").str[0]
            tmp["post"] = tmp["group"].str.split(".").str[-1]
            ptbl = pl.from_pandas(tmp)

            pre_groups = ptbl.group_by("pre", maintain_order=True).agg(
                pl.col("simple_id")
            )
            for row in pre_groups.iter_rows(named=True):
                sid = np.asarray(row["simple_id"], dtype=int)
                pre_ret[row["pre"]] = agg_func(y[..., sid], -1)

            post_groups = ptbl.group_by("post", maintain_order=True).agg(
                pl.col("simple_id")
            )
            for row in post_groups.iter_rows(named=True):
                sid = np.asarray(row["simple_id"], dtype=int)
                post_ret[row["post"]] = agg_func(y[..., sid], -1)
        else:
            tmp["pre"] = tmp["group"].apply(lambda x: x.split(".")[0])
            tmp["post"] = tmp["group"].apply(lambda x: x.split(".")[-1])
            for pre, group in tmp.groupby("pre", dropna=True):
                pre_ret[pre] = agg_func(y[..., group.simple_id], -1)
            for post, group in tmp.groupby("post", dropna=True):
                post_ret[post] = agg_func(y[..., group.simple_id], -1)
        return pre_ret, post_ret
    if mode == "all_innervated":
        assert (
            connections is not None
        ), "connections must be provided for all_innervated mode"
        tmp = connections[["pre_simple_id", "post_simple_id", "neuropil"]]
        tmp = tmp[
            (tmp["pre_simple_id"] < y.shape[-1]) & (tmp["post_simple_id"] < y.shape[-1])
        ]
        pre_ret: dict = {}
        post_ret: dict = {}

        if use_polars:
            ptbl = pl.from_pandas(tmp)
            groups = ptbl.group_by("neuropil", maintain_order=True).agg(
                pl.col("pre_simple_id"), pl.col("post_simple_id")
            )
            for row in groups.iter_rows(named=True):
                neuropil = row["neuropil"]
                pre_ids = np.asarray(row["pre_simple_id"], dtype=int)
                post_ids = np.asarray(row["post_simple_id"], dtype=int)
                pre_ret[neuropil] = agg_func(y[..., pre_ids], -1)
                post_ret[neuropil] = agg_func(y[..., post_ids], -1)
        else:
            for neuropil, group in tmp.groupby("neuropil", dropna=True):
                pre_ret[neuropil] = agg_func(y[..., group.pre_simple_id.to_numpy()], -1)
                post_ret[neuropil] = agg_func(
                    y[..., group.post_simple_id.to_numpy()], -1
                )
        return pre_ret, post_ret

    raise ValueError(
        "Invalid `mode` for `agg_by_neuropil`. "
        "Expected one of: {'top_innervated', 'all_innervated'}."
    )

def agg_conn(
    y,
    conn: pd.DataFrame,
    conn_weight: scipy.sparse.sparray | None = None,
    neurons: pd.DataFrame | None = None,
    mode: Literal["neuropil", "neuron"] = "neuron",
    neuron_type_column: str = "cell_type",
    agg: Literal["mean", "sum", "std"] = "mean",
):
    """Aggregate connectivity weights by neuropil or neuron-type pairs."""
    if conn_weight is not None:
        conn_weight = conn_weight.tocoo()
        conn = conn.merge(
            pd.DataFrame(
                {
                    "pre_simple_id": conn_weight.row,
                    "post_simple_id": conn_weight.col,
                    "weight": conn_weight.data,
                }
            ),
            how="left",
            on=["pre_simple_id", "post_simple_id"],
        )
    if mode == "neuropil":
        return conn.groupby("neuropil")["weight"].agg(agg)
    if mode == "neuron":
        assert neurons is not None, "neurons must be provided for neuron mode"
        conn = conn.merge(
            neurons[["simple_id", neuron_type_column]].rename(
                columns={
                    "simple_id": "pre_simple_id",
                    neuron_type_column: f"pre_{neuron_type_column}",
                }
            ),
            how="left",
            on="pre_simple_id",
        )
        conn = conn.merge(
            neurons[["simple_id", neuron_type_column]].rename(
                columns={
                    "simple_id": "post_simple_id",
                    neuron_type_column: f"post_{neuron_type_column}",
                }
            ),
            how="left",
            on="post_simple_id",
        )
        return conn.groupby(
            [f"pre_{neuron_type_column}", f"post_{neuron_type_column}"]
        )["weight"].agg(agg)

    raise ValueError(
        "Invalid `mode` for `agg_conn`. Expected one of: {'neuropil', 'neuron'}."
    )

def branching_ratio(
    all_counts: np.ndarray | list[np.ndarray] | list[list[float]] | str,
    k_max: int = 40,
    *,
    maxslopes: int | None = None,
    scatterpoints: bool = False,
    eps: float = 1e-12,
    ar1_fallback: bool = True,
) -> dict:
    """Estimate branching ratio via subsampling-invariant MR fitting.

    Citation:
    Wilting, J., & Priesemann, V. (2018). Inferring collective dynamical states
    from widely unobserved systems. Nature Communications, 9(1), 2325.
    https://doi.org/10.1038/s41467-018-04725-4

    Args:
        all_counts: One trial, stacked trials, list of trials, or path input.
        k_max: Maximum lag used for the MR fit.
        maxslopes: Legacy synonym for ``k_max``.
        scatterpoints: Keep centered lagged ``x`` and ``y`` samples in result.
        eps: Small stabilizer for divisions and weights.
        ar1_fallback: Use AR(1) fallback when MR fit is ill-posed.

    Returns:
        Dictionary with ``branching_ratio``, ``naive_branching_ratio``, ``k``,
        ``r_k``, ``stderr`` and fit metadata.

    Raises:
        ValueError: If input data are invalid or too short.
        RuntimeError: If MR and fallback estimation both fail.
    """
    counts_list = input_handler(all_counts)

    if maxslopes is not None:
        k_max = maxslopes
    k_max = int(max(2, k_max))

    shortest = min(len(c) for c in counts_list)
    max_possible = min(k_max, shortest - 1)
    if max_possible < 1:
        raise ValueError("Time series too short for branching-ratio estimation")

    k, r_k_raw, stderr_raw, data_length, mean_activity, xs, ys = get_slopes(
        counts_list,
        max_possible + 1,
        scatterpoints=scatterpoints,
        eps=eps,
    )

    r_k = np.asarray(r_k_raw, dtype=np.float64)
    stderr = np.asarray(stderr_raw, dtype=np.float64)

    valid = np.isfinite(r_k) & np.isfinite(stderr) & (r_k > 0) & (stderr > 0)
    if np.sum(valid) < 2:
        if ar1_fallback:
            result = _ar1_fallback(counts_list, data_length, mean_activity, eps)
            if scatterpoints:
                result["xs"] = xs
                result["ys"] = ys
            return result
        raise RuntimeError("Insufficient valid r_k values for MR estimation")

    k_valid = k[valid].astype(np.float64)
    r_valid = np.clip(r_k[valid], 1e-12, 1e12)
    stderr_valid = stderr[valid]

    y = np.log(r_valid)
    w = 1.0 / (stderr_valid + eps)
    w = w / np.mean(w)

    slope, intercept = np.polyfit(k_valid, y, 1, w=w)
    m_hat = float(np.exp(slope))
    a_fit = float(np.exp(intercept))

    if not np.isfinite(m_hat) or m_hat <= 0:
        raise RuntimeError("Fitted branching ratio is non-positive or invalid")

    result = {
        "branching_ratio": m_hat,
        "a_fit": a_fit,
        "autocorrelationtime": float(-1.0 / np.log(m_hat))
        if (m_hat > 0 and m_hat != 1.0)
        else np.inf,
        "naive_branching_ratio": float(r_k[0]) if len(r_k) > 0 else np.nan,
        "k": k,
        "r_k": r_k,
        "stderr": stderr,
        "fit_slope": float(slope),
        "fit_intercept": float(intercept),
        "fit_points_used": int(np.sum(valid)),
        "data_length": data_length,
        "mean_activity": mean_activity,
    }
    if scatterpoints:
        result["xs"] = xs
        result["ys"] = ys
    return result

def build_group_frame(
    values: TensorLike,
    neurons_df: pd.DataFrame,
    group_by: str,
    *,
    simple_id_col: str = "simple_id",
    value_name: str = "value",
    dropna: bool = True,
) -> pd.DataFrame:
    """Convert neuron-aligned values to a tidy frame for grouped analyses.

    Args:
        values: Array/tensor shaped `[N]` or `[..., N]` where the last axis is
            neuron. All leading dimensions are flattened into independent
            samples (e.g., trials, conditions, or time points).
        neurons_df: DataFrame containing at least `simple_id_col` and `group_by`.
        group_by: Column in `neurons_df` used as grouping key.
        simple_id_col: Column mapping rows in `neurons_df` to neuron index in
            `values`.
        value_name: Name for the output value column.
        dropna: Drop missing values in group/value columns when `True`.
    """
    y = _to_numpy(values)
    if y.ndim < 1:
        raise ValueError("`values` must have at least one dimension.")

    if simple_id_col not in neurons_df.columns:
        raise ValueError(f"Missing `{simple_id_col}` in `neurons_df`.")
    if group_by not in neurons_df.columns:
        raise ValueError(f"Missing `{group_by}` in `neurons_df`.")

    metadata = neurons_df.loc[:, [simple_id_col, group_by]].copy()
    if dropna:
        metadata = metadata.dropna(subset=[group_by])
    if metadata.empty:
        raise ValueError("No neuron metadata available after filtering.")

    if metadata[simple_id_col].duplicated().any():
        raise ValueError(f"`{simple_id_col}` must be unique in `neurons_df`.")

    try:
        simple_ids = pd.to_numeric(metadata[simple_id_col], errors="raise").to_numpy(
            dtype=np.int64
        )
    except Exception as exc:
        raise ValueError(f"`{simple_id_col}` must be numeric.") from exc

    n_neurons = y.shape[-1]
    out_of_range = (simple_ids < 0) | (simple_ids >= n_neurons)
    if np.any(out_of_range):
        bad_ids = simple_ids[out_of_range]
        raise ValueError(
            f"Found `{simple_id_col}` outside [0, {n_neurons - 1}]: {bad_ids.tolist()}"
        )

    selected = y[..., simple_ids]
    n_samples = int(np.prod(selected.shape[:-1], dtype=np.int64))
    n_samples = max(1, n_samples)

    flattened = selected.reshape(n_samples, len(simple_ids))
    group_labels = metadata[group_by].to_numpy()

    frame = pd.DataFrame(
        {
            group_by: np.repeat(group_labels, n_samples),
            value_name: flattened.T.reshape(-1),
        }
    )

    if dropna:
        frame = frame.dropna(subset=[value_name])
    if frame.empty:
        raise ValueError("No values available after filtering.")

    return frame

def compute_ie_ratio(
    excitatory_mat: sparray,
    inhibitory_mat: sparray,
    excitatory_neuron_only: bool = True,
    neurons: Optional[pd.DataFrame] = None,
    warn_strict: bool = True,
) -> tuple[float, np.ndarray]:
    """Compute inhibitory/excitatory ratio per neuron and whole-brain mean."""
    if excitatory_neuron_only:
        assert neurons is not None

    sum_inhibitory = inhibitory_mat.sum(axis=0).astype(float)
    sum_excitatory = excitatory_mat.sum(axis=0).astype(float)

    # neurons that have zero inputs from both e and i connections are likely input,
    # so don't warn on them
    input_indices = (sum_excitatory == 0) & (sum_inhibitory == 0)

    sum_excitatory[sum_excitatory == 0] = np.nan
    if excitatory_neuron_only:
        sum_excitatory[neurons[neurons.EI == "E"].simple_id.to_numpy()] = np.nan

    ie_ratios = sum_inhibitory / sum_excitatory

    if warn_strict:
        nan_indices = (np.isnan(ie_ratios) & ~input_indices).nonzero()[0]
        if nan_indices.size > 0:
            print(f"Warning: IE ratio contains NaN values at indices {nan_indices}")

        inf_indices = np.isinf(ie_ratios).nonzero()[0]
        if inf_indices.size > 0:
            print(f"Warning: IE ratio contains Inf values at indices {inf_indices}")

    ie_ratios = np.where(np.isinf(ie_ratios), np.nan, ie_ratios)
    ie_ratio_whole_brain: float = np.nanmean(ie_ratios[ie_ratios != 0])

    return ie_ratio_whole_brain, ie_ratios

def compute_log_hist(data, bins=1000, edge_pos: Literal["mid", "sep"] = "mid"):
    """Compute histogram with logarithmically-spaced bins.

    Useful for visualizing heavy-tailed distributions like synaptic weights
    or degree distributions.

    Args:
        data: Input data array (must be positive).
        bins: Number of histogram bins.
        edge_pos: Position to return for bin edges.
            "mid": Return bin centers (midpoints).
            "sep": Return bin separators (edges).

    Returns:
        Tuple of (hist, bin_edges) where hist is the count array and
        bin_edges are the positions (centers or edges based on edge_pos).

    Raises:
        ValueError: If data contains non-positive values (required for log scale).
    """
    bin_edges = np.logspace(np.log10(np.min(data)), np.log10(np.max(data)), num=bins)
    hist, edges = np.histogram(data, bins=bin_edges)
    if edge_pos == "mid":
        bin_edges = 0.5 * (edges[:-1] + edges[1:])
    return hist, bin_edges

def compute_raster(sp_matrix: np.ndarray, times: np.ndarray):
    """Get spike raster plot which displays the spiking activity of a group of
    neurons over time."""
    times = np.asarray(times)
    elements = np.where(sp_matrix > 0.0)
    index = elements[1]
    time = times[elements[0]]
    return index, time

def compute_spectrum(y, dt, nperseg=None):
    from scipy.signal import welch

    freqs, Y_mag = welch(y, fs=1 / dt, nperseg=nperseg, axis=0)
    return freqs, Y_mag

@use_stats(value_key="cv_temporal", dim=1)
def cv_temporal(
    spike_data: np.ndarray | torch.Tensor,
    dt_ms: float = 1.0,
    window: int = 100,
    step: int = 1,
    batch_axis: tuple[int, ...] | int | None = None,
    dtype: np.dtype | torch.dtype | None = None,
):
    """Compute CV in sliding temporal windows.

    Calculates the coefficient of variation of ISIs within sliding windows
    over time, giving a time-resolved measure of spike train irregularity.

    Args:
        spike_data: Spike train array of shape [T, ...]. First dimension is time.
        dt_ms: Time step in milliseconds.
        window: Size of the sliding window in time steps.
        step: Step size between consecutive windows.
        batch_axis: Axes to aggregate across (e.g., (1, 2) for trials).
        dtype: Data type for output arrays. If None, uses float64 for NumPy
            and float32 for Torch.

    Returns:
        cv_temporal: CV values for each window. Shape: [n_windows, ...]
            where n_windows = (T - window) // step + 1.
        info: Dictionary with window boundaries.
    """
    if isinstance(batch_axis, int):
        batch_axis = (batch_axis,)

    T = spike_data.shape[0]
    n_windows = (T - window) // step + 1

    # Determine output dtype
    if isinstance(spike_data, torch.Tensor):
        out_dtype = dtype if dtype is not None else torch.float32
        cv_values = torch.full(
            (n_windows,) + spike_data.shape[1:],
            float("nan"),
            dtype=out_dtype,
            device=spike_data.device,
        )
    else:
        out_dtype = dtype if dtype is not None else float
        cv_values = np.full(
            (n_windows,) + spike_data.shape[1:], np.nan, dtype=out_dtype
        )

    for i in range(n_windows):
        start = i * step
        end = start + window
        if end > T:
            break

        window_data = spike_data[start:end]

        cv_window, _ = _cv(window_data, dt_ms, batch_axis, dtype)
        cv_values[i] = cv_window

    window_starts = np.arange(n_windows) * step * dt_ms
    window_ends = window_starts + window * dt_ms

    info = {
        "window": window,
        "step": step,
        "window_starts_ms": window_starts,
        "window_ends_ms": window_ends,
    }
    return cv_values, info

def describe_array(array: np.ndarray):
    """Print descriptive statistics for a 1D array.

    Displays mean, median, std, min, max, and quartiles.

    Args:
        array: 1D NumPy array to describe.

    Example:
        >>> describe_array(np.random.randn(100))
        Mean: 0.05
        Median: 0.12
        ...
    """
    mean = np.mean(array)
    median = np.median(array)
    std_dev = np.std(array)
    min_val = np.min(array)
    max_val = np.max(array)
    q25 = np.percentile(array, 25)
    q50 = np.percentile(array, 50)  # This is the same as the median
    q75 = np.percentile(array, 75)

    print(f"Mean: {mean}")
    print(f"Median: {median}")
    print(f"Standard Deviation: {std_dev}")
    print(f"Min: {min_val}")
    print(f"Max: {max_val}")
    print(f"25th Percentile (Q1): {q25}")
    print(f"50th Percentile (Q2/Median): {q50}")
    print(f"75th Percentile (Q3): {q75}")