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 choose_current_clamp_sweeps(
    sweep_records: Sequence[dict[str, Any]],
    *,
    preferred_stimuli: Sequence[str] = ("long square",),
) -> list[dict[str, Any]]:
    """Select current-clamp sweeps appropriate for somatic fitting."""
    matched = []
    for sweep in sweep_records:
        stimulus_name = (
            sweep.get("stimulus_name")
            or sweep.get("ephys_stimulus", {}).get("description")
            or ""
        )
        stimulus_units = sweep.get("stimulus_units") or ""
        clamp_mode = sweep.get("clamp_mode") or sweep.get("stimulus_type_name") or ""

        is_current_clamp = _matches_any(stimulus_units, ["pa", "amp"]) or (
            clamp_mode and not _matches_any(clamp_mode, ["voltage clamp"])
        )
        if not is_current_clamp:
            continue
        if preferred_stimuli and not _matches_any(stimulus_name, preferred_stimuli):
            continue
        matched.append(sweep)
    return matched

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