Analysis — Dynamic Tools¶

def calculate_kaplan_yorke_dimension(lyapunov_spectrum: np.ndarray):
    """Calculate the Kaplan-Yorke Dimension (D_KY), also known as the Lyapunov
    Dimension.

    Formula: D_KY = k + sum(lambda_i for i=1 to k) / |lambda_{k+1}|
    where k is the max index such that the sum of the first k exponents is non-negative.

    Args:
        lyapunov_spectrum (np.ndarray): Array of Lyapunov exponents, sorted in
        descending order.

    Returns:
        float: The Kaplan-Yorke dimension. Returns 0 if the system is stable
            (all lambda < 0). Returns the number of exponents if the sum of all
            is positive (unbounded/hyperchaos).
    """
    # Ensure sorted descending
    ls = np.sort(lyapunov_spectrum)[::-1]

    n = len(ls)

    # Calculate cumulative sums
    cum_sum = np.cumsum(ls)

    # Find k: max index such that sum >= 0
    # We look for the last index where cum_sum >= 0
    positive_sums = np.where(cum_sum >= 0)[0]

    if len(positive_sums) == 0:
        # All cumulative sums are negative.
        # This usually means the first exponent is negative (stable fixed point).
        return 0.0

    k = positive_sums[-1]

    # Check if k is the last element (sum of all is positive)
    if k == n - 1:
        return float(n)

    # Apply formula
    # Note: indices are 0-based in Python, so k corresponds to the (k+1)-th
    # element in 1-based math notation.
    # The formula uses 1-based k.
    # Let's map carefully:
    # Python index k is the index of the last element included in the sum.
    # So we have summed ls[0]...ls[k].
    # The next element is ls[k+1].
    # The integer part of dimension is (k + 1).

    sum_lambda = cum_sum[k]
    lambda_next = ls[k + 1]

    if lambda_next == 0:
        # Avoid division by zero, though theoretically lambda_{k+1} should be
        # negative here.
        return float(k + 1)

    d_ky = (k + 1) + sum_lambda / abs(lambda_next)

    return d_ky

def calculate_structural_eigenvalue_outliers(
    weight_matrix: np.ndarray, spectral_radius: float = None
):
    """Analyze the eigenvalues of the weight matrix to identify structural
    outliers.

    According to the circular law, eigenvalues of a random matrix are distributed
    within a disk of radius R. Outliers outside this radius indicate structural
    enforcement of specific oscillatory modes (stable dynamics) rather than
    random chaos.

    Args:
        weight_matrix (np.ndarray): The connectivity weight matrix (N x N).
        spectral_radius (float, optional): The theoretical spectral radius of the
            random component. If None, it is estimated as std(W) * sqrt(N).

    Returns:
        dict: Dictionary containing:
            - 'eigenvalues': All eigenvalues.
            - 'outliers': Eigenvalues outside the spectral radius.
            - 'outlier_count': Number of outliers.
            - 'spectral_radius': The radius used for thresholding.
    """
    # Ensure numpy array
    W = np.array(weight_matrix)
    N = W.shape[0]

    if W.shape[0] != W.shape[1]:
        raise ValueError("Weight matrix must be square.")

    # Compute eigenvalues
    eigenvalues = np.linalg.eigvals(W)

    # Determine spectral radius if not provided
    if spectral_radius is None:
        # Estimate radius based on random matrix theory
        # For entries with variance sigma^2/N, radius is sigma.
        # Here we have entries with variance var(W).
        # If W_ij ~ N(0, sigma^2), then radius R = sigma * sqrt(N).
        # std(W) corresponds to sigma.
        sigma = np.std(W)
        spectral_radius = sigma * np.sqrt(N)

    # Identify outliers
    magnitudes = np.abs(eigenvalues)
    outlier_indices = np.where(magnitudes > spectral_radius)[0]
    outliers = eigenvalues[outlier_indices]

    return {
        "eigenvalues": eigenvalues,
        # True Spectral Radius
        "max_eigenvalue": np.max(magnitudes) if len(magnitudes) > 0 else 0.0,
        "outliers": outliers,
        "outlier_count": len(outliers),
        "spectral_radius": spectral_radius,  # Bulk Radius (Threshold)
    }

def calculate_gain_stability_sensitivity(
    model, dataloader, g_values=None, dt=1.0, device="cuda"
) -> float:
    """Calculate the Gain-Stability Sensitivity (Susceptibility) slope.

    Definition: The slope of the curve of the Maximum Lyapunov Exponent (lambda_max)
    as a function of global synaptic gain scaling (g).

    Args:
        model: The Brain model.
        dataloader: DataLoader providing input.
        g_values: List of gain scaling factors. Default np.linspace(0.5, 5.0, 10).
        dt: Simulation time step.
        device: Device to run on.

    Returns:
        tuple: (slope, intercept, g_values, lambda_values)
            - slope: The slope of lambda_max vs g.
            - intercept: The intercept of the fit.
            - g_values: The gain values used.
            - lambda_values: The computed max Lyapunov exponents.
    """
    from model import functional, init

    if g_values is None:
        g_values = np.linspace(0.5, 5.0, 10)

    # Access linear layer
    # Assuming model is Brain, model.brain is RecurrentNN, model.brain.synapse
    # is Synapse model.brain.synapse.linear is the layer
    try:
        linear_layer = model.brain.synapse.linear
    except AttributeError:
        print("Could not find linear layer at model.brain.synapse.linear")
        return 0.0

    original_magnitude = linear_layer.magnitude.data.clone()

    lambda_values = []

    model.eval()
    model.to(device)

    # Get one batch of input
    try:
        batch = next(iter(dataloader))
    except StopIteration:
        print("Dataloader is empty.")
        return 0.0

    inputs = batch["input"]
    # inputs: (Batch, Time, ...) -> (Time, Batch, ...)
    inputs = inputs.transpose(0, 1).to(device)

    # We can use the first sample in the batch.
    input_sample = inputs[:, 0:1, ...]  # Keep batch dim 1

    for g in g_values:
        # Scale weights
        linear_layer.magnitude.data = original_magnitude * g

        # Reset state
        functional.reset_net(model, device=device)
        init.uniform_v_(model.brain.neuron, set_reset_value=True, batch_size=1)

        # Run
        with torch.no_grad():
            _, brain_out = model(input_sample)
            spikes = brain_out["neuron"]["spike"]  # (Time, Batch, Neurons)

        # Convert to rate
        # spikes: (Time, 1, Neurons) -> (Time, Neurons)
        spikes_sq = spikes.squeeze(1)

        # Continuous rate
        rates = get_continuous_spiking_rate(spikes_sq, dt=dt)

        # Mean population rate for LE calculation
        mean_rate = rates.mean(axis=1)

        # Compute LE
        try:
            le = compute_max_lyapunov_exponent(mean_rate)
        except Exception as e:
            print(f"Error computing LE for g={g}: {e}")
            le = 0.0  # Or NaN?

        lambda_values.append(le)

    # Restore weights
    linear_layer.magnitude.data = original_magnitude

    # Calculate slope
    # Fit line: lambda = slope * g + intercept
    # Handle potential NaNs or Infs
    valid_indices = np.isfinite(lambda_values)
    if np.sum(valid_indices) < 2:
        return 0.0

    g_valid = g_values[valid_indices]
    lambda_valid = np.array(lambda_values)[valid_indices]

    slope, intercept = np.polyfit(g_valid, lambda_valid, 1)

    return slope, intercept, g_values, np.array(lambda_values)

def calculate_lyapunov_exponent(spike_train: torch.Tensor, dt: float = 0.1) -> float:
    """Calculate the maximum Lyapunov exponent for a given spike train.

    Args:
        spike_train (torch.Tensor): The spike train data. Shape (time_steps,
        num_neurons).
        dt (float): Time bin size in milliseconds. Default is 0.1 ms.

    Returns:
        float: The maximum Lyapunov exponent.
    """
    # Ensure spike_train is a 2D tensor
    if spike_train.ndim != 2:
        raise ValueError(
            "spike_train must be a 2D tensor with shape (time_steps, num_neurons)"
        )

    # 1. Calculate the continuous spiking rate using a Gaussian kernel (smooth
    # the spike train) This is effectively a form of kernel density estimation.
    # We use a small bandwidth, as the original dynamics should be captured at a
    # fine timescale.
    bandwidth = 5.0  # in ms, this may need adjustment
    continuous_rate = get_continuous_spiking_rate(spike_train, dt, bandwidth)

    # 2. Calculate the Lyapunov exponent using the continuous rate
    # We use the largest Lyapunov exponent as the measure of chaos/complexity.
    lyapunov_exponent = compute_max_lyapunov_exponent(continuous_rate, dt)

    return lyapunov_exponent

def calculate_pcist(
    response: torch.Tensor, baseline: torch.Tensor, threshold_factor: float = 3.0
) -> float:
    """Calculate the Perturbational Complexity Index based on State Transitions
    (PCIst).

    Definition: A measure of the spatiotemporal complexity of the network's
    response to a specific perturbation.

    Steps:
    1. Perturb (assumed done, input is response).
    2. Measure (input is response matrix).
    3. Decompose: Perform PCA on the response matrix.
    4. Recurrence: Calculate state transitions on the principal components.
    5. Sum significant state transitions weighted by the component's
    Signal-to-Noise ratio.

    Args:
        response (torch.Tensor): The network response to perturbation. Shape
        (time_steps, num_neurons).
        baseline (torch.Tensor): The baseline
        activity before perturbation. Shape (time_steps_base, num_neurons).
        threshold_factor (float): Factor of baseline std dev to define
        significant state excursion. Default 3.0.

    Returns:
        float: The PCIst score.
    """
    # Ensure inputs are float tensors
    if not isinstance(response, torch.Tensor):
        response = torch.tensor(response, dtype=torch.float32)
    if not isinstance(baseline, torch.Tensor):
        baseline = torch.tensor(baseline, dtype=torch.float32)

    response = response.float()
    baseline = baseline.float()

    # Handle batch dimension: if 3D, calculate mean PCIst over batch or raise
    # error? For simplicity, if 3D, we assume (batch, time, neurons) and
    # calculate average PCIst.
    if response.ndim == 3:
        batch_size = response.shape[0]
        pcist_values = []
        for i in range(batch_size):
            # Handle corresponding baseline
            b_sample = baseline[i] if baseline.ndim == 3 else baseline
            pcist_values.append(
                calculate_pcist(response[i], b_sample, threshold_factor)
            )
        return sum(pcist_values) / len(pcist_values)

    # 1. Center data based on baseline mean
    mean_base = baseline.mean(dim=0)
    response_centered = response - mean_base
    baseline_centered = baseline - mean_base

    # 2. PCA on Response
    # We use SVD for PCA: X = U S V^T. Principal components (scores) are X V = U S.
    # response_centered shape: (T, N)
    try:
        # full_matrices=False ensures we get min(T, N) components
        U, S, Vh = torch.linalg.svd(response_centered, full_matrices=False)
    except RuntimeError:
        # Fallback for singular matrices or convergence issues
        return 0.0

    V = Vh.T  # (N, K)

    # Project data onto Principal Components
    # Scores shape: (T, K)
    scores_response = torch.matmul(response_centered, V)
    scores_baseline = torch.matmul(baseline_centered, V)

    # 3. Calculate SNR for each component
    # SNR = Variance(Response) / Variance(Baseline)
    # Add epsilon to avoid division by zero
    var_response = scores_response.var(dim=0)
    var_baseline = scores_baseline.var(dim=0)

    epsilon = 1e-9
    snr = var_response / (var_baseline + epsilon)

    # 4. Calculate State Transitions
    # A state transition is defined as crossing a threshold defined by baseline noise.
    # Threshold for component k: threshold_factor * std(baseline_k)

    std_baseline = scores_baseline.std(dim=0)
    thresholds = threshold_factor * std_baseline  # (K,)

    # Binarize response: 1 if |score| > threshold, else 0
    # We are looking for "significant excursions"
    # Shape: (T, K)
    active_states = (torch.abs(scores_response) > thresholds.unsqueeze(0)).float()

    # Count transitions: change from 0 to 1 or 1 to 0
    # diff along time dimension
    transitions = torch.abs(active_states[1:] - active_states[:-1])

    # Sum transitions for each component
    num_transitions = transitions.sum(dim=0)  # (K,)

    # 5. Weighted Sum "Sum significant state transitions weighted by the
    # component's Signal-to-Noise ratio." We might want to filter components
    # with SNR < 1? The text doesn't strictly say so, but "weighted by SNR"
    # implies low SNR components contribute little. However, if SNR is very
    # high, it dominates. Let's follow the instruction literally.

    pcist_score = (num_transitions * snr).sum()

    return pcist_score.item()

def calculate_ra(spike_initial: torch.Tensor, spike_final: torch.Tensor) -> float:
    """Calculate Representation Alignment (RA) using spike data.

    RA = Trace(G_final * G_initial) / (||G_final|| * ||G_initial||)
    where G = S * S^T (Gram matrix of spike activity)

    If inputs are 3D tensors, they are assumed to be (batch_size, time_steps,
    num_neurons) and will be averaged over the time dimension (dim=1) to obtain
    firing rates.

    Args:
        spike_initial (torch.Tensor): Initial spike activity. Shape (batch_size,
        num_neurons) or (batch_size, time_steps, num_neurons). spike_final
        (torch.Tensor): Final spike activity. Shape (batch_size, num_neurons) or
        (batch_size, time_steps, num_neurons).

    Returns:
        float: The Representation Alignment (RA) score.
               Low RA -> Rich Regime (Radical restructuring)
               High RA -> Lazy Regime (Little change in internal structure)
    """
    # Ensure inputs are float tensors
    if not isinstance(spike_initial, torch.Tensor):
        spike_initial = torch.tensor(spike_initial, dtype=torch.float32)
    if not isinstance(spike_final, torch.Tensor):
        spike_final = torch.tensor(spike_final, dtype=torch.float32)

    spike_initial = spike_initial.float()
    spike_final = spike_final.float()

    # Handle 3D input: (batch, time, neurons) -> (batch, neurons)
    if spike_initial.ndim == 3:
        spike_initial = spike_initial.mean(dim=1)
    if spike_final.ndim == 3:
        spike_final = spike_final.mean(dim=1)

    # 1. Compute Gram matrix G = S * S^T
    # Shape: (batch_size, batch_size)
    g_initial = torch.matmul(spike_initial, spike_initial.T)
    g_final = torch.matmul(spike_final, spike_final.T)

    # 2. Calculate Trace(G_final * G_initial)
    # Trace(A @ B) = sum(element-wise product of A and B^T)
    # Since G is symmetric, G^T = G, so this is sum(G_final * G_initial)
    product = torch.matmul(g_final, g_initial)
    numerator = torch.trace(product)

    # 3. Calculate norms ||G||
    # Assuming Frobenius norm as is standard for matrix alignment
    norm_initial = torch.norm(g_initial, p="fro")
    norm_final = torch.norm(g_final, p="fro")

    # 4. Calculate RA
    if norm_initial == 0 or norm_final == 0:
        return 0.0  # Avoid division by zero

    ra = numerator / (norm_final * norm_initial)

    return ra.item()

def compute_max_lyapunov_exponent(time_series, emb_dim=6, lag=1, tau=1):
    """Compute the largest Lyapunov exponent of a given time series using the
    nolds library.

    Parameters:
    - time_series: A 1D numpy array representing the time series data.

    Returns:
    - lyapunov_exponent: The estimated largest Lyapunov exponent.
    """
    lyapunov_exponent = nolds.lyap_r(time_series, emb_dim=emb_dim, lag=lag, tau=tau)
    return lyapunov_exponent

def get_continuous_spiking_rate(spikes, dt, sigma=20.0):
    """Convert discrete spike trains into continuous firing rates using
    Gaussian smoothing.

    Args:
        spikes (np.ndarray or torch.Tensor): Spike matrix of shape (time_steps,
        n_neurons).
        dt (float): Simulation time step in ms.
        sigma (float): Standard deviation of the Gaussian kernel in ms. Default 20ms.

    Returns:
        np.ndarray: Continuous firing rate traces of shape (time_steps, n_neurons).
    """
    if isinstance(spikes, torch.Tensor):
        spikes = spikes.detach().cpu().numpy()

    # Convert sigma from ms to bins
    sigma_bins = sigma / dt

    # Apply Gaussian filter along the time axis (axis 0)
    rates = gaussian_filter1d(spikes.astype(float), sigma=sigma_bins, axis=0)

    return rates

def _fit_distribution(data):
    """Helper to fit power law distribution using powerlaw package."""
    if len(data) < 10:
        return np.nan, None
    try:
        # discrete=True because sizes/durations are counts (integers)
        fit = powerlaw.Fit(data, discrete=True, verbose=False)
        return fit.alpha, fit
    except Exception as e:
        warnings.warn(f"Failed to fit power law distribution: {e}")
        return np.nan, None

def _fit_scaling(x, y):
    """Helper to fit power law scaling y = a * x^gamma using curve_fit."""
    if len(x) < 3:
        return np.nan, None

    try:
        # Initial guess: a=1, gamma=1.5
        popt, pcov = curve_fit(_power_law_func, x, y, p0=[1, 1.5], maxfev=2000)
        gamma = popt[1]

        # Calculate R^2
        residuals = y - _power_law_func(x, *popt)
        ss_res = np.sum(residuals**2)
        ss_tot = np.sum((y - np.mean(y)) ** 2)
        r_squared = 1 - (ss_res / ss_tot)

        stats = {"r_squared": r_squared, "popt": popt, "pcov": pcov}
        return gamma, stats
    except Exception as e:
        warnings.warn(f"Failed to fit scaling relation: {e}")
        return np.nan, None

def _power_law_func(x, a, gamma):
    return a * np.power(x, gamma)

def calculate_dfa(spike_train: np.ndarray, bin_size: int = 1):
    """Calculate Detrended Fluctuation Analysis (DFA) exponent alpha.

    Meaning of alpha:
    - 0.5: White noise (no memory)
    - 0.5 < alpha < 1.0: Long-range memory (fractal structure)
    - 1.0: 1/f noise (Pink noise)
    - 1.5: Brownian motion (Random walk)

    Args:
        spike_train (np.ndarray): Binary spike matrix of shape (time_steps, n_neurons).
        bin_size (int): Width of time bin in number of time steps.

    Returns:
        float: The DFA exponent alpha.
    """
    if not HAS_NOLDS:
        raise ImportError(
            "nolds package is required for DFA analysis. "
            "Install with: pip install nolds"
        )

    # Ensure input is numpy array
    spike_train = np.array(spike_train)

    # 1. Calculate population activity (sum spikes across neurons)
    if spike_train.ndim == 2:
        population_activity = np.sum(spike_train, axis=1)
    else:
        population_activity = spike_train

    # 2. Binning
    if bin_size > 1:
        n_bins = len(population_activity) // bin_size
        population_activity = population_activity[: n_bins * bin_size]
        population_activity = population_activity.reshape(-1, bin_size).sum(axis=1)

    # 3. Calculate DFA using nolds
    # nolds.dfa expects the time series (it performs integration internally)
    try:
        alpha = nolds.dfa(population_activity)
        return alpha
    except Exception as e:
        warnings.warn(f"Failed to calculate DFA: {e}")
        return np.nan

def compute_avalanche_statistics(spike_train: np.ndarray, bin_size: int = 1):
    """Calculate avalanche size (S) and duration (T) distributions and their
    power-law exponents.

    Definition: An avalanche is defined as a continuous sequence of time bins
    (width bin_size) containing at least one spike, flanked by empty bins.

    Args:
        spike_train (np.ndarray): Binary spike matrix of shape (time_steps, n_neurons).
        bin_size (int): Width of time bin in number of time steps.

    Returns:
        dict: Dictionary containing:
            - 'tau': Power-law exponent for avalanche size distribution P(S) ~
              S^-tau
            - 'alpha': Power-law exponent for avalanche duration distribution
              P(T) ~ T^-alpha
            - 'gamma': Power-law exponent for average size vs duration <S>(T) ~
              T^gamma
            - 'gamma_pred': Predicted gamma based on tau and alpha:
              (alpha-1)/(tau-1)
            - 'CCC': Criticality Consistency Coefficient: 1 - |gamma -
              gamma_pred| / gamma
            - 'sizes': List of avalanche sizes
            - 'durations': List of avalanche durations
            - 'avg_size_by_duration': Tuple (unique_durations, mean_sizes)
            - 'fit_S': powerlaw.Fit object for sizes
            - 'fit_T': powerlaw.Fit object for durations
    """
    # Ensure input is numpy array
    spike_train = np.array(spike_train)

    # Check dimensions. We expect (Time, Neurons).
    if spike_train.ndim != 2:
        raise ValueError("spike_train must be a 2D matrix (time_steps, n_neurons)")

    # 1. Calculate population activity (sum spikes across neurons)
    population_activity = np.sum(spike_train, axis=1)  # Shape: (T,)

    # 2. Binning
    if bin_size > 1:
        n_bins = len(population_activity) // bin_size
        # Truncate to multiple of bin_size
        population_activity = population_activity[: n_bins * bin_size]
        # Reshape and sum
        population_activity = population_activity.reshape(-1, bin_size).sum(axis=1)

    # 3. Identify avalanches
    # Active bins are those with > 0 spikes
    is_active = population_activity > 0

    # Find continuous sequences of active bins
    # Pad with False to detect start/end at boundaries
    padded_active = np.concatenate(([False], is_active, [False]))
    diff = np.diff(padded_active.astype(int))

    starts = np.where(diff == 1)[0]
    ends = np.where(diff == -1)[0]

    sizes = []
    durations = []

    for start, end in zip(starts, ends):
        # segment from start to end (exclusive)
        segment = population_activity[start:end]

        # Size (S): Total number of spikes in the avalanche
        s = np.sum(segment)

        # Duration (T): Number of time bins the avalanche lasts
        t = len(segment)  # equivalent to end - start

        sizes.append(s)
        durations.append(t)

    sizes = np.array(sizes)
    durations = np.array(durations)

    results = {
        "sizes": sizes,
        "durations": durations,
        "tau": np.nan,
        "alpha": np.nan,
        "gamma": np.nan,
        "gamma_pred": np.nan,
        "CCC": np.nan,
        "fit_S": None,
        "fit_T": None,
    }

    if len(sizes) < 10:
        warnings.warn(
            f"Not enough avalanches to fit power law. Found {len(sizes)} avalanches."
        )
        return results

    # 4. Fit power laws using MLE (powerlaw package)
    results["tau"], results["fit_S"] = _fit_distribution(sizes)
    results["alpha"], results["fit_T"] = _fit_distribution(durations)

    # 5. Average Size vs. Duration Scaling (<S>(T) ~ T^gamma)
    if len(durations) > 0:
        # Use bincount for fast grouping by integer duration
        counts = np.bincount(durations)
        sum_sizes = np.bincount(durations, weights=sizes)

        # Filter out durations that didn't occur
        mask = counts > 0
        unique_durations = np.arange(len(counts))[mask]
        mean_sizes = sum_sizes[mask] / counts[mask]

        results["avg_size_by_duration"] = (unique_durations, mean_sizes)

        # Fit scaling relation using curve_fit (non-linear least squares)
        results["gamma"], results["gamma_stats"] = _fit_scaling(
            unique_durations, mean_sizes
        )

    # 6. Calculate Criticality Consistency Coefficient (CCC)
    # gamma_pred = (alpha - 1) / (tau - 1)
    # CCC = 1 - |gamma_obs - gamma_pred| / gamma_obs
    if (
        not np.isnan(results["tau"])
        and not np.isnan(results["alpha"])
        and not np.isnan(results["gamma"])
    ):
        try:
            if results["tau"] != 1:
                gamma_pred = (results["alpha"] - 1) / (results["tau"] - 1)
                results["gamma_pred"] = gamma_pred

                if results["gamma"] != 0:
                    ccc = 1 - abs(results["gamma"] - gamma_pred) / results["gamma"]
                    results["CCC"] = ccc
        except Exception as e:
            warnings.warn(f"Failed to calculate CCC: {e}")

    return results

def _compute_eci(
    I_e: torch.Tensor | np.ndarray,
    I_i: torch.Tensor | np.ndarray,
    *,
    I_ext: torch.Tensor | np.ndarray | None = None,
    batch_axis: tuple[int, ...] | int | None = None,
    dtype: torch.dtype | np.dtype | None = None,
) -> torch.Tensor | np.ndarray:
    if isinstance(batch_axis, int):
        batch_axis = (batch_axis,)

    # Check if all inputs are zero - return ones in that case
    if (I_e == 0).all() and (I_i == 0).all() and (I_ext is None or (I_ext == 0).all()):
        # Determine which axes are aggregated to compute output shape
        if batch_axis is not None:
            agg_axes = {0} | set(batch_axis)  # time (0) + batch axes
        else:
            agg_axes = set(range(I_e.ndim - 1))  # all except last (neurons)
        output_shape = tuple(I_e.shape[i] for i in range(I_e.ndim) if i not in agg_axes)
        if len(output_shape) == 0:
            output_shape = (1,)
        if isinstance(I_e, torch.Tensor):
            return torch.ones(output_shape, dtype=I_e.dtype, device=I_e.device)
        return np.ones(output_shape, dtype=I_e.dtype)

    # Handle I_ext by splitting into excitatory and inhibitory parts
    if I_ext is not None:
        if isinstance(I_e, torch.Tensor):
            I_ext_pos = torch.clamp(I_ext, min=0)
            I_ext_neg = torch.clamp(I_ext, max=0)
        else:
            I_ext_pos = np.clip(I_ext, 0, None)
            I_ext_neg = np.clip(I_ext, None, 0)
        I_e_eff = I_e + I_ext_pos
        I_i_eff = I_i + I_ext_neg
    else:
        I_e_eff = I_e
        I_i_eff = I_i

    # Compute recurrent current
    I_rec = I_e_eff + I_i_eff

    # Determine axes/dims for aggregation
    if batch_axis is not None:
        agg_dims = (0,) + tuple(batch_axis)
    else:
        agg_dims = tuple(range(I_e.ndim - 1))

    if isinstance(I_e, torch.Tensor):
        numer = torch.abs(I_rec).mean(dim=agg_dims, dtype=dtype)
        denom = (torch.abs(I_e_eff) + torch.abs(I_i_eff)).mean(
            dim=agg_dims, dtype=dtype
        )
        denom = denom + torch.finfo(I_e.dtype).eps
        return numer / denom
    else:
        numer = np.abs(I_rec).mean(axis=agg_dims, dtype=dtype)
        denom = (np.abs(I_e_eff) + np.abs(I_i_eff)).mean(axis=agg_dims, dtype=dtype)
        denom = denom + np.finfo(I_e.dtype).eps
        return numer / denom

def _compute_lag_correlation(
    x: torch.Tensor | np.ndarray,
    y: torch.Tensor | np.ndarray,
    *,
    dt: float = 1.0,
    max_lag_ms: float = 30.0,
    batch_axis: tuple[int, ...] | int | None = None,
    use_fft: bool = True,
    dtype: torch.dtype | np.dtype | None = None,
):
    assert max_lag_ms >= dt
    if isinstance(x, torch.Tensor):
        y = torch.as_tensor(y, dtype=x.dtype, device=x.device)
    else:
        y = y.numpy(force=True) if isinstance(y, torch.Tensor) else np.asarray(y)

    is_torch = isinstance(x, torch.Tensor)
    T = x.shape[0]
    n_neurons = x.shape[-1]
    lag_bins = int(max_lag_ms / dt)
    max_lag = min(lag_bins, T - 1)

    if isinstance(batch_axis, int):
        batch_axis = (batch_axis,)

    # Common reshape logic (works for both torch and numpy)
    x_3d = x.reshape(T, -1, n_neurons)
    y_3d = y.reshape(T, -1, n_neurons)
    flat_x = x_3d.reshape(T, -1)
    flat_y = y_3d.reshape(T, -1)

    if use_fft or T > 100:
        corr_flat = _cross_correlation_fft(flat_x, flat_y, max_lag, dtype)
    else:
        corr_flat = _cross_correlation_direct(flat_x, flat_y, max_lag, dtype)

    n_lags = corr_flat.shape[0]
    corr_3d = corr_flat.reshape(n_lags, -1, n_neurons)

    if is_torch:
        if batch_axis is not None:
            corr_agg = corr_3d.mean(dim=1, dtype=dtype)
        else:
            corr_agg = corr_3d.reshape(n_lags, -1)

        peak_corr, best_lag_idx = torch.max(corr_agg, dim=0)
        max_lag_actual = min(lag_bins, T - 1)
        best_lags = best_lag_idx - max_lag_actual
        best_lag_ms = best_lags * dt

        info = {
            "corr_over_lags": corr_agg,
            "lag_values_ms": torch.arange(
                -max_lag_actual, max_lag_actual + 1, device=x.device
            )
            * dt,
        }
        return peak_corr, best_lag_ms, info
    else:
        if batch_axis is not None:
            corr_agg = corr_3d.mean(axis=1)
        else:
            corr_agg = corr_3d.reshape(n_lags, -1)

        best_lag_idx = np.argmax(corr_agg, axis=0)
        peak_corr = corr_agg[best_lag_idx, np.arange(corr_agg.shape[1])]
        max_lag_actual = min(lag_bins, T - 1)
        best_lags = best_lag_idx - max_lag_actual
        best_lag_ms = best_lags * dt

        info = {
            "corr_over_lags": corr_agg,
            "lag_values_ms": np.arange(-max_lag_actual, max_lag_actual + 1) * dt,
        }
        return peak_corr, best_lag_ms, info

def _cross_correlation_direct(
    x: torch.Tensor | np.ndarray,
    y: torch.Tensor | np.ndarray,
    max_lag: int,
    dtype: torch.dtype | np.dtype | None = None,
) -> torch.Tensor | np.ndarray:
    """Direct cross-correlation (simpler for short signals)."""
    is_torch = isinstance(x, torch.Tensor)
    T, N = x.shape

    if is_torch:
        device = x.device
        # Normalize
        x_mean = x.mean(dim=0, keepdim=True)
        y_mean = y.mean(dim=0, keepdim=True)
        x_std = x.std(dim=0, keepdim=True) + torch.finfo(x.dtype).eps
        y_std = y.std(dim=0, keepdim=True) + torch.finfo(x.dtype).eps
        x_norm = (x - x_mean) / x_std
        y_norm = (y - y_mean) / y_std
        # Compute correlations for each lag
        max_lag_actual = min(max_lag, T - 1)
        n_lags = 2 * max_lag_actual + 1
        corr = torch.zeros(n_lags, N, device=device, dtype=x.dtype)
        for i, lag in enumerate(range(-max_lag_actual, max_lag_actual + 1)):
            if lag < 0:
                c = (x_norm[:lag] * y_norm[-lag:]).mean(dim=0, dtype=dtype)
            elif lag > 0:
                c = (x_norm[lag:] * y_norm[:-lag]).mean(dim=0, dtype=dtype)
            else:
                c = (x_norm * y_norm).mean(dim=0, dtype=dtype)
            corr[i, :] = c
        return corr
    else:
        # Normalize
        x_norm = (x - x.mean(axis=0)) / (x.std(axis=0) + np.finfo(x.dtype).eps)
        y_norm = (y - y.mean(axis=0)) / (y.std(axis=0) + np.finfo(x.dtype).eps)
        # Compute correlations for each lag
        n_lags = 2 * max_lag + 1
        corr = np.zeros((n_lags, N))
        for i, lag in enumerate(range(-max_lag, max_lag + 1)):
            if lag < 0:
                c = (x_norm[:lag] * y_norm[-lag:]).mean(axis=0, dtype=dtype)
            elif lag > 0:
                c = (x_norm[lag:] * y_norm[:-lag]).mean(axis=0, dtype=dtype)
            else:
                c = (x_norm * y_norm).mean(axis=0, dtype=dtype)
            corr[i, :] = c
        return corr

def _cross_correlation_fft(
    x: torch.Tensor | np.ndarray,
    y: torch.Tensor | np.ndarray,
    max_lag: int,
    dtype: torch.dtype | np.dtype | None = None,
) -> torch.Tensor | np.ndarray:
    """FFT-based cross-correlation."""
    is_torch = isinstance(x, torch.Tensor)
    T, N = x.shape
    n_fft = 2 * T

    if is_torch:
        # Demean and compute FFT
        x_demean = x - x.mean(dim=0, keepdim=True)
        y_demean = y - y.mean(dim=0, keepdim=True)
        X = torch.fft.rfft(x_demean.float(), n=n_fft, dim=0)
        Y = torch.fft.rfft(y_demean.float(), n=n_fft, dim=0)
        cross_spec = X * Y.conj()
        corr_full = torch.fft.irfft(cross_spec, n=n_fft, dim=0)
        # Normalize
        x_std = x.std(dim=0, keepdim=True) + torch.finfo(x.dtype).eps
        y_std = y.std(dim=0, keepdim=True) + torch.finfo(x.dtype).eps
        corr_norm = corr_full / (x_std * y_std * T)
        # Extract valid lags
        max_lag_actual = min(max_lag, T - 1)
        neg_lags = corr_norm[-max_lag_actual:, :]
        pos_lags = corr_norm[: max_lag_actual + 1, :]
        return torch.cat([neg_lags, pos_lags], dim=0)
    else:
        # Demean and compute FFT
        x_demean = x - x.mean(axis=0, keepdims=True, dtype=dtype)
        y_demean = y - y.mean(axis=0, keepdims=True, dtype=dtype)
        X = np.fft.rfft(x_demean, n=n_fft, axis=0)
        Y = np.fft.rfft(y_demean, n=n_fft, axis=0)
        cross_spec = X * np.conj(Y)
        corr_full = np.fft.irfft(cross_spec, n=n_fft, axis=0)
        # Normalize
        x_std = x.std(axis=0, keepdims=True, dtype=dtype) + np.finfo(x.dtype).eps
        y_std = y.std(axis=0, keepdims=True, dtype=dtype) + np.finfo(x.dtype).eps
        corr_norm = corr_full / (x_std * y_std * T)
        # Extract valid lags
        neg_lags = corr_norm[-max_lag:, :]
        pos_lags = corr_norm[: max_lag + 1, :]
        return np.concatenate([neg_lags, pos_lags], axis=0)