Pipeline functions

def prepare_timeseries(data, tr=0.72, bandpass=(0.01, 0.1),
                       n_regions=None, mapping=None):
    """Time x regions; filter, select regions, then average grouped signals."""
    X = np.asarray(data, dtype=float)
    if X.ndim != 2 or X.shape[0] < 4 or X.shape[1] < 2:
        raise ValueError("Input must be a time x regions matrix.")
    if not np.isfinite(X).all():
        raise ValueError("Input contains NaN or infinite values.")
    if tr <= 0:
        raise ValueError("TR must be positive.")
    if bandpass is not None:
        low, high = bandpass
        if not 0 < low < high < 0.5 / tr:
            raise ValueError("Bandpass must lie between zero and Nyquist.")
        b, a = signal.butter(2, [low, high], btype="bandpass", fs=1 / tr)
        X = signal.filtfilt(b, a, X, axis=0)
    if n_regions is not None:
        if not 2 <= n_regions <= X.shape[1]:
            raise ValueError("Invalid number of selected regions.")
        X = X[:, :n_regions]
    if mapping is not None:
        mapping = np.asarray(mapping).reshape(-1)
        labels = np.unique(mapping)
        if (len(mapping) != X.shape[1] or not np.isfinite(mapping).all()
                or not np.array_equal(labels, np.arange(1, len(labels) + 1))):
            raise ValueError("Mapping must contain consecutive integer labels starting at 1.")
        X = np.column_stack([X[:, mapping == k].mean(axis=1) for k in labels])
    return X

def empirical_moments(X, lag=1):
    """Reproduce calc_Corr0/calc_Corrlag in my_module.py exactly.

    Q0 uses X[:-1] even when lag differs from 1. Each segment is centered
    separately and standardized with population standard deviations (ddof=0),
    while the product is divided by segment length minus one.
    Thus Q0's diagonal is (T-1)/(T-2), not exactly one.
    Q1[i,j] relates earlier region i to later region j.
    """
    X = np.asarray(X, dtype=float)
    if X.ndim != 2 or not np.isfinite(X).all():
        raise ValueError("Expected finite time x regions data.")
    if not isinstance(lag, (int, np.integer)) or not 1 <= lag < X.shape[0] - 1:
        raise ValueError("lag must be a positive integer smaller than T-1.")
    def standardized(A):
        sd = A.std(axis=0)
        if np.any(sd == 0):
            raise ValueError("Constant time series encountered.")
        return (A - A.mean(axis=0)) / sd
    Z0 = standardized(X[:-1])
    Zpast = standardized(X[:-lag])
    Zfuture = standardized(X[lag:])
    Q0 = Z0.T @ Z0 / (len(Z0) - 1)
    Q1 = Zpast.T @ Zfuture / (len(Zpast) - 1)
    return Q0, Q1

def fit_tau_grid(data, taus, tr=0.72, bandpass=(0.01, 0.1), n_regions=100,
                mapping=None, max_iter=100, eta_C=0.005):
    """Fit independently at each tau; do not refit the selected estimate."""
    X = prepare_timeseries(data, tr, bandpass, n_regions, mapping)
    Q0, Q1 = empirical_moments(X, lag=1)
    taus = np.asarray(taus, float)
    if taus.ndim != 1 or len(taus) < 4 or not np.isfinite(taus).all() or np.any(taus <= 0) or np.any(np.diff(taus) <= 0):
        raise ValueError("Use at least four strictly increasing positive finite taus.")
    if max_iter <= 12:
        raise ValueError("max_iter must exceed 12.")
    Js = np.full((len(taus), len(Q0), len(Q0)), np.nan)
    Sigmas = np.full_like(Js, np.nan)
    iterations = np.zeros(len(taus), int)
    stable = np.zeros(len(taus), bool)
    distances = np.full(len(taus), np.nan)
    errors = [''] * len(taus)
    for i, tau in enumerate(taus):
        try:
            model = MOU()
            model.fit_LO(Q0, Q1, lag=1, prior=0.0, tau=float(tau),
                         eta_C=eta_C, max_iter=max_iter)
            J, Sigma = np.asarray(model.J), np.asarray(model.Sigma)
            if Sigma.ndim == 1:
                Sigma = np.diag(Sigma)
            iterations[i] = int(model.d_fit['iterations'])
            distances[i] = float(model.d_fit['distance'])
            Js[i], Sigmas[i] = J, Sigma
            if np.isfinite(J).all() and np.isfinite(Sigma).all() and np.isfinite(distances[i]):
                stable[i] = np.linalg.eigvals(J).real.max() < 0
        except (ValueError, FloatingPointError, np.linalg.LinAlgError, UnboundLocalError, OverflowError) as exc:
            errors[i] = type(exc).__name__ + ': ' + str(exc)
    # Return the fitted grid without choosing tau or computing asymmetries.
    return dict(taus=taus, J_grid=Js, Sigma_grid=Sigmas, iterations=iterations,
                stable=stable, distance=distances, errors=np.asarray(errors), Q0=Q0, Q1=Q1,
                tr=np.asarray(tr), max_iter=np.asarray(max_iter))

def select_precision_reference(fitted):
    """Choose the precision-reference minimum in the convergent interval."""
    selected = dict(fitted)
    taus, matrices = selected['taus'], selected['J_grid']
    Q0, Q1 = selected['Q0'], selected['Q1']
    offdiag = ~np.eye(len(Q0), dtype=bool)
    def pearson(first, second):
        if (np.isfinite(first).all() and np.isfinite(second).all()
                and first.std() > 0 and second.std() > 0):
            return float(np.corrcoef(first, second)[0, 1])
        return np.nan

    # In raw pyMOU orientation the reference is -Q0^-1 Sigma.
    # Transposing both matrices gives the equivalent paper convention.
    scores = np.full(len(taus), np.nan)
    for i, (J, Sigma) in enumerate(zip(matrices, selected['Sigma_grid'])):
        if np.isfinite(Sigma).all():
            reference = -0.5 * np.linalg.solve(Q0, Sigma)
            scores[i] = pearson(J[offdiag], reference[offdiag])

    iterations = selected['iterations']
    eligible = (selected['stable'] & (iterations > 0)
                & (iterations < int(selected['max_iter'])) & np.isfinite(scores)
                & np.array([np.any(J[offdiag] != 0) for J in matrices]))
    search, initial = np.zeros(len(taus), bool), np.zeros(len(taus), bool)
    asym_scores = np.full(len(taus), np.nan)
    indices = np.flatnonzero(eligible)
    if len(indices):
        first = last = int(indices[0])
        while last + 1 < len(taus) and eligible[last + 1]:
            last += 1
        search[first:last + 1] = True
        # Require antisymmetric correlation below 0.9 before selecting tau.
        asym_Q1 = (Q1 - Q1.T) / 2
        asym_scores = np.array([pearson(((J - J.T) / 2)[offdiag], asym_Q1[offdiag])
                                for J in matrices])
        above_threshold = search & (~np.isfinite(asym_scores) | (asym_scores >= 0.9))
        excluded_indices = np.flatnonzero(above_threshold)
        if len(excluded_indices):
            initial = search & (np.arange(len(taus)) <= excluded_indices[-1])
        search &= ~initial
    remaining = np.flatnonzero(search)
    outliers = np.zeros(len(taus), bool)
    if len(remaining):
        outliers = (np.arange(len(taus)) <= remaining[0]) & np.isfinite(scores) & (scores < 0.25)
    search &= ~outliers
    candidates = np.flatnonzero(search)
    # Select an interior local minimum, not the lowest endpoint of the curve.
    # Both immediate neighbours must also be admissible grid points.
    local_minima = np.zeros(len(taus), dtype=bool)
    if len(taus) >= 3:
        local_minima[1:-1] = (search[:-2] & search[1:-1] & search[2:]
                             & (scores[1:-1] <= scores[:-2])
                             & (scores[1:-1] <= scores[2:])
                             & ((scores[1:-1] < scores[:-2])
                                | (scores[1:-1] < scores[2:])))
    minima = np.flatnonzero(local_minima)
    best = int(minima[np.argmin(scores[minima])]) if len(minima) else -1
    boundary = False
    selected.update(corr_precision=scores, eligible=eligible, search_mask=search,
                    local_minima=local_minima, tau_identified=np.asarray(best >= 0),
                    corr_asym_J_Q1=asym_scores, asym_corr_threshold=np.asarray(0.9),
                    excluded_initial=initial, red_outliers=outliers,
                    best_index=np.asarray(best), tau_opt=np.asarray(taus[best] if best >= 0 else np.nan),
                    boundary_minimum=np.asarray(boundary),
                    selection=np.asarray('Interior local minimum of off-diagonal corr(J_raw, -Q0^-1 Sigma); asymmetry correlation < 0.9'))
    selected.pop('boundary_maximum', None)
    empty = np.full_like(Q0, np.nan)
    selected['J_raw'] = matrices[best].copy() if best >= 0 else empty
    selected['J_target_source'] = selected['J_raw'].T.copy()
    selected['Sigma'] = selected['Sigma_grid'][best].copy() if best >= 0 else empty.copy()

    # Keep pairwise correlations so the similarity plot can be reloaded without
    # saving every EC matrix fitted along the tau grid.
    vectors = matrices[:, offdiag]
    valid = np.isfinite(vectors).all(axis=1) & (vectors.std(axis=1) > 0)
    correlations = np.full((len(taus), len(taus)), np.nan)
    if valid.any():
        correlations[np.ix_(valid, valid)] = np.atleast_2d(np.corrcoef(vectors[valid]))
    selected['ec_correlations_all'] = correlations
    return selected