How much recording do you need, and what does more of it buy?

Jansen-Rit columns are much harder than phase oscillators: each column’s alpha rhythm is driven mostly by its own dynamics, so a neighbour’s influence is a small perturbation on top.

Recall climbs with recording length and then saturates, but precision drops, because with enough data the indirect path a -> b -> c shows up as a -> c. Those extra detections are not errors, information really does flow that way, they just are not edges of the network you wanted. So there is an optimum recording length, which is not what most people expect.

Run with:

uv run python -m examples.neural_mass_network [--quick]

alpha check: peak 10.74 Hz, 100% of power in 8-13 Hz
simulating and testing   30s of recording ...
/home/runner/work/pyPTE/pyPTE/examples/neural_mass_network.py:81: RuntimeWarning: 40 surrogates cannot produce a p-value below 0.0244, but Benjamini-Hochberg over 10 channel pairs needs 0.005 to accept an isolated edge. Unless the effect is widespread, expect nothing to survive. Use at least 199 surrogates, or test at the cluster level with cluster_permutation().
  result = surrogate_test(signal, n_surrogates=n_surrogates, seed=0)
simulating and testing   80s of recording ...
simulating and testing  160s of recording ...

results
   seconds   found  false+   recall  precision
        30       0       0     0.00       0.00
        80       5       5     1.00       0.50
       160       5       9     1.00       0.36

verified: columns oscillate at 10.7 Hz with 100% of power in the alpha band,
          and 30s of data is not enough to find coupling that 160s does find.

          Watch precision rather than recall: the extra detections at
          long recordings are indirect paths, not errors. For network
          reconstruction there is an optimum recording length.

import argparse

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import ListedColormap
from matplotlib.patches import Patch
from scipy.signal import welch

from examples.models import kuramoto, neural_mass_model
from examples.utils import style
from examples.utils.plotting import FIGURES, plot_matrix, save, use_house_style
from pyPTE import surrogate_test

FS = 250.0
COUPLING_STRENGTH = 30.0
DELAY = 0.020  # about a fifth of an alpha cycle, where direction is clearest
N_COLUMNS = 5


def score(result, truth: np.ndarray) -> dict:
    """Precision and recall of the significant edges against ground truth."""
    off_diagonal = ~np.eye(truth.shape[0], dtype=bool)
    found = result.significant & off_diagonal
    true_edges = truth & off_diagonal

    true_positives = int((found & true_edges).sum())
    false_positives = int((found & ~true_edges).sum())

    return {
        "true_positives": true_positives,
        "false_positives": false_positives,
        "recall": true_positives / max(int(true_edges.sum()), 1),
        "precision": true_positives / max(true_positives + false_positives, 1),
        "significant": found,
    }


def main(quick: bool = False) -> None:
    use_house_style()
    n_surrogates = 40 if quick else 100
    durations = [30, 80, 160] if quick else [30, 60, 90, 120, 160, 200]

    coupling = kuramoto.ring_with_shortcut(N_COLUMNS, strength=COUPLING_STRENGTH)
    truth = coupling > 0

    # -- confirm the columns really are producing an alpha rhythm ----------
    reference = neural_mass_model.simulate(
        coupling, n_samples=int(60 * FS), delay=DELAY, seed=1
    )
    freqs, power = welch(reference[0], fs=FS, nperseg=1024)
    peak = float(freqs[np.argmax(power)])
    alpha_fraction = float(power[(freqs >= 8) & (freqs < 13)].sum() / power.sum())
    print(f"alpha check: peak {peak:.2f} Hz, {alpha_fraction:.0%} of power in 8-13 Hz")

    # -- how does detection change with recording length? ------------------
    rows = []
    for seconds in durations:
        print(f"simulating and testing {seconds:>4}s of recording ...")
        signal = neural_mass_model.simulate(
            coupling, n_samples=int(seconds * FS), delay=DELAY, seed=1
        )
        result = surrogate_test(signal, n_surrogates=n_surrogates, seed=0)
        rows.append({"seconds": seconds, "result": result, **score(result, truth)})

    fig, axes = plt.subplots(2, 2, figsize=(10.5, 7.5), constrained_layout=True)

    window = slice(0, int(3 * FS))
    for column in range(min(3, N_COLUMNS)):
        axes[0, 0].plot(
            np.arange(int(3 * FS)) / FS,
            reference[column, window],
            lw=0.9,
            label=f"column {column}",
        )
    axes[0, 0].set(
        xlabel="time (s)",
        ylabel="pyramidal potential (mV)",
        title="Jansen-Rit columns produce an alpha rhythm",
    )
    axes[0, 0].legend(fontsize=7, ncols=3)

    axes[0, 1].semilogy(freqs, power, lw=1.2, color=style.BLUE)
    axes[0, 1].axvspan(8, 13, color=style.OBS, alpha=0.15)
    axes[0, 1].set(
        xlim=(0, 40),
        xlabel="frequency (Hz)",
        ylabel="power",
        title=f"Spectrum: peak {peak:.1f} Hz",
    )

    seconds = [row["seconds"] for row in rows]
    axes[1, 0].plot(
        seconds,
        [row["recall"] for row in rows],
        "o-",
        color=style.BLUE,
        label="recall (edges found)",
    )
    axes[1, 0].plot(
        seconds,
        [row["precision"] for row in rows],
        "s-",
        color=style.OBS,
        label="precision (of those, correct)",
    )
    axes[1, 0].set(
        xlabel="recording length (s)",
        ylabel="score",
        ylim=(-0.05, 1.05),
        title="More data finds every edge, then adds wrong ones",
    )
    axes[1, 0].legend(fontsize=8, loc="lower left")

    # 0 correct rejection, 1 false positive, 2 missed edge, 3 true positive
    longest = rows[-1]
    outcome = 2 * truth.astype(int) + longest["significant"].astype(int)
    categories = ListedColormap([style.TINT, style.OBS, style.AMBER, style.BLUE])
    plot_matrix(
        axes[1, 1],
        outcome.astype(float),
        title=(
            f"Ground truth vs findings at {longest['seconds']}s\n"
            f"precision {longest['precision']:.2f}, recall {longest['recall']:.2f}"
        ),
        cmap=categories,
        vmin=-0.5,
        vmax=3.5,
        colorbar=False,
        mask_diagonal=True,
    )
    axes[1, 1].legend(
        handles=[
            Patch(color=style.BLUE, label="true edge, found"),
            Patch(color=style.AMBER, label="true edge, missed"),
            Patch(color=style.OBS, label="no edge, but flagged"),
            Patch(color=style.TINT, label="no edge, not flagged"),
        ],
        fontsize=7,
        loc="upper left",
        bbox_to_anchor=(1.02, 1.0),
    )

    save(fig, FIGURES / "neural_mass_network.png")

    print("\nresults")
    print(
        f"  {'seconds':>8} {'found':>7} {'false+':>7} {'recall':>8} {'precision':>10}"
    )
    for row in rows:
        print(
            f"  {row['seconds']:>8} {row['true_positives']:>7} "
            f"{row['false_positives']:>7} {row['recall']:>8.2f} "
            f"{row['precision']:>10.2f}"
        )

    best_recall = max(row["recall"] for row in rows)
    shortest = rows[0]

    assert 8.0 <= peak <= 13.0, (
        f"columns should oscillate in the alpha band, got {peak}"
    )
    assert alpha_fraction > 0.5, "most of the power should sit in the alpha band"
    assert shortest["recall"] < best_recall, (
        "a short recording should miss edges that a longer one finds"
    )
    assert best_recall >= 0.8, "a long enough recording should find nearly every edge"

    print(
        f"\nverified: columns oscillate at {peak:.1f} Hz with "
        f"{alpha_fraction:.0%} of power in the alpha band,"
        f"\n          and {shortest['seconds']}s of data is not enough to find "
        f"coupling that {durations[-1]}s does find."
        "\n\n          Watch precision rather than recall: the extra detections at"
        "\n          long recordings are indirect paths, not errors. For network"
        "\n          reconstruction there is an optimum recording length."
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--quick", action="store_true", help="fewer surrogates")
    main(**vars(parser.parse_args()))