diff --git a/docs/python/matlab-vs-python.md b/docs/python/matlab-vs-python.md index 24d750de..c596a024 100644 --- a/docs/python/matlab-vs-python.md +++ b/docs/python/matlab-vs-python.md @@ -48,6 +48,11 @@ Non-negative matrix factorization (NMF) uses `scikit-learn`'s coordinate descent solver rather than MATLAB's Alternating Least Squares — a genuinely different algorithm, not just a different random seed. Treat `num_nnmf_components` and related outputs as approximate. + +The variance-explained curve (`nnmf_var_explained`, `nnmf_residuals`) also +stops early by default: once it passes 95%, or rises by less than +`nmf_plateau_gain` (0.005) over 10 ranks. MATLAB fits every rank up to the +number of active nodes; set `nmf_plateau_gain` to 0 to do the same. ::: :::: diff --git a/packaging/meanap.spec b/packaging/meanap.spec index 6c049352..042548cd 100644 --- a/packaging/meanap.spec +++ b/packaging/meanap.spec @@ -37,6 +37,9 @@ hiddenimports = ( + collect_submodules("spikeinterface.postprocessing") + collect_submodules("spikeinterface.qualitymetrics") + ["mountainsort5", "hdbscan", "oasis", "oasis.functions", "certifi"] + # numba loads its parallel threading layer (omppool / workqueue / tbbpool) + # by name the first time a ``parallel=True`` function runs. + + collect_submodules("numba.np.ufunc") ) excludes = [ diff --git a/python/PIPELINE_PORT_STATUS.md b/python/PIPELINE_PORT_STATUS.md index 114f623c..a84b895b 100644 --- a/python/PIPELINE_PORT_STATUS.md +++ b/python/PIPELINE_PORT_STATUS.md @@ -763,7 +763,59 @@ a negative control), and the 5G figures go through `test_stats_report.py`. - CAT-NAP runs name their columns `Grp`/`DIV`; MATLAB's `doStats` expects `eGrp`/`AgeDiv`. `dataset.META_COLUMNS` knows both. -### How to verify changes +#### Hundreds of sorted units — 7x end to end (2026-10-24) + +Spike-sorted recordings make every node-count-dependent cost bite. Measured on +the MaxOne retinal organoid in `local/test-spike-sorted-data/` (306 units, +919k spikes, 126 s; steps 1-4, lags 10/15/25 ms, 200 shuffles, full figures, +one recording, 16-core box): + +| | Before | After | +|---|---|---| +| Step 3 | 1494s | **121s** | +| Step 4 | 1749s | **332s** | +| **Total** | **3258s** | **463s** | + +(The "before" run was under cProfile; its step 3 is ~2.2s x 603 STTC matrices +unprofiled, so the comparison holds.) Adjacency matrices and both network CSVs +are identical before and after; the NMF arrays agree to 1e-14 (see below). + +Where it went, and what changed: + +| Cost (before) | Change | After | +|---|---|---| +| STTC pair sweep, 1489s | `get_sttc`'s pair loop is one numba call, rows across cores (`prange`, rows paired from both ends to balance). The work is O(units x total spikes) per matrix, so compiling alone gained nothing; the cores did. | 2.2s -> 0.22s per matrix | +| NMF, 933s | fits run on threads (sklearn's CD update releases the GIL), each with a seed drawn up front in the old order; the per-rank sweep stops at 95% variance explained or a plateau (`nmf_plateau_gain`, default 0.005 per 10 ranks; 0 = every rank, as MATLAB). Here it stops at rank 228 of 306. | 26s | +| `participation_coef_norm`, 379s | its 100 null models are built on threads. Each model's random numbers depend only on the original network's edge count, so they are all drawn first, in the sequential order (`_draw_null_plan`), then the rewiring (`randmio_und_signed`, now `nogil`) runs in parallel. Same models and same generator state as the loop — this is what the 2026-08-26 note above said could not be done without changing results. | 33s | +| `randmio_und_v2`/`latmio_und_v2`, <1s | compiled, fed batches from the caller's generator so the draws keep their order; identical output. Negligible here (they were cheap at 306 nodes), done for completeness. | ~0s | + +What is left in step 4 is **figures** (~250s, of which `plot_network` is +~190s for 72 network plots at 306 nodes); its computation is 86s. Express mode +skips them. + +**Results.** Everything above is exact except one thing: NMF now holds BLAS +to one thread for every fit (threadpoolctl). Before, it ran with whatever BLAS +share the machine gave it, and a multi-threaded BLAS sums in a +thread-count-dependent order, so NMF already differed between machines and +worker counts in the last bits. It no longer does; against an old run its +arrays move by ~1e-14. `num_nnmf_components` did not move here, but on another +recording such noise could flip it by one, as it could have between machines +before. + +**Threads.** `parallel.compute_threads()` gives a process its share: a +`map_recordings` worker's pinned BLAS share, otherwise every physical core. +STTC sets numba's thread count from it; NMF and the null models size their +thread pools from it. If numba cannot load a threading layer (a frozen build +missing one), `get_sttc` warns once and runs the same kernel serially; the +app smoke test now exercises the parallel kernel, and the PyInstaller spec +collects `numba.np.ufunc`. + +Checked by `python/test_pipeline_fast_paths.py` [3]-[6] (each rewrite against +the form it replaced, at one thread and several, including generator state +afterwards), plus the step-3/step-4 MATLAB parity tests, CAT-NAP, sorted import +and stim-network tests. + +## How to verify changes ``` uv run python python/test_stats.py # 47 checks, ~2 min — the analyses diff --git a/python/test_gui_advanced.py b/python/test_gui_advanced.py index b739765b..6b079447 100644 --- a/python/test_gui_advanced.py +++ b/python/test_gui_advanced.py @@ -87,8 +87,9 @@ def check(name: str, condition: bool, detail: str = "") -> None: app.processEvents() sections = panel.findChildren(AdvancedSection) -# One under Probabilistic thresholding, one under Node and edge inclusion. -check("connectivity has two advanced sections", len(sections) == 2, str(len(sections))) +# One each under Probabilistic thresholding, Node and edge inclusion, and +# Network metrics. +check("connectivity has three advanced sections", len(sections) == 3, str(len(sections))) # Loaded into a collapsed panel, read straight back out. loaded = Params(prob_thresh_tail=0.02, diff --git a/python/test_pipeline_fast_paths.py b/python/test_pipeline_fast_paths.py index 9e59fd30..4695ed36 100644 --- a/python/test_pipeline_fast_paths.py +++ b/python/test_pipeline_fast_paths.py @@ -4,11 +4,12 @@ uv run python python/test_pipeline_fast_paths.py -Three routines were rewritten for speed after profiling showed a 270-cell +Routines rewritten for speed — first after profiling showed a 270-cell calcium recording spending ~960s on a single lag (``null_model_und_sign`` and -``efficiency_wei_local`` between them accounting for ~90% of it). Each rewrite -is supposed to be *exactly* equivalent, not merely close, so each is checked -here against the slower form it replaced: +``efficiency_wei_local`` between them accounting for ~90% of it), then for +spike-sorted recordings with hundreds of units. Each rewrite is supposed to be +*exactly* equivalent, not merely close, so each is checked here against the +slower form it replaced: 1. ``null_models._stable_ranks`` — an ``argpartition`` that reads only the ranks the caller asks for, instead of sorting every remaining edge. It @@ -19,8 +20,23 @@ 2. ``network_metrics.distance_wei`` — a numba-compiled Dijkstra, checked against the NumPy implementation it now shadows (which is still the fallback when numba is unavailable). -3. ``sttc.get_sttc`` — hoists ``run_T`` out of the pairwise loop, so it is - checked against a direct sweep of the unchanged ``sttc_pair``. +3. ``sttc.get_sttc`` — hoists ``run_T`` out of the pairwise loop and runs + the whole sweep as one compiled, row-parallel call, so it is checked + against a direct sweep of the unchanged ``sttc_pair``, at one thread and + at several. +4. ``null_models.randmio_und_v2`` / ``latmio_und_v2`` — the swap loop is + compiled and fed random numbers in batches from the caller's generator. + Checked against the generator-based version it replaced (kept below as + the reference): same networks *and* the same generator state afterwards, + so everything drawn later in a run is unchanged too. +5. ``nmf.cal_nmf`` — fits run on threads, each with a seed drawn up front in + the old order. Checked against the old serial loops at several thread + counts, and the early-stopping sweep against the full one. +6. ``null_models.map_null_models_und_sign`` — participation_coef_norm's 100 + null models, built on threads from random numbers drawn up front. Checked + against null models that draw straight from the generator as they go (the + pre-threading path, still used without numba): same networks, same + generator state afterwards. """ from __future__ import annotations @@ -34,7 +50,10 @@ sys.path.insert(0, str(REPO_ROOT / "src")) from meanap.pipeline import network_metrics as nm # noqa: E402 -from meanap.pipeline.null_models import _stable_ranks # noqa: E402 +from meanap.pipeline import nmf # noqa: E402 +from meanap.pipeline.null_models import ( # noqa: E402 + _null_from_plan, _stable_ranks, latmio_und_v2, map_null_models_und_sign, randmio_und_v2, +) from meanap.pipeline.sttc import get_sttc, sttc_pair # noqa: E402 FAILURES: list[str] = [] @@ -155,27 +174,312 @@ def random_graph(n: int, density: float, seed: int) -> np.ndarray: print() -print("[3] get_sttc — hoisted run_T vs the per-pair sttc_pair") +print("[3] get_sttc — compiled parallel sweep vs the per-pair sttc_pair") -for n, n_spikes, duration in ((30, 40, 120.0), (60, 8, 300.0), (12, 1, 60.0)): +import numba # noqa: E402 + + +def sttc_reference(trains: dict, n: int, lag_ms: float, duration: float) -> np.ndarray: + want = np.full((n, n), np.nan) + for i in range(n): + for j in range(i + 1, n): + c = sttc_pair(trains[i], trains[j], lag_ms / 1000.0, 0.0, duration) + want[i, j] = c + want[j, i] = c + want[want < 0] = 0.0 + want[np.isnan(want)] = 0.0 + return want + + +# Odd and even n, since the parallel sweep takes rows in pairs from both ends +# and the middle row of an odd n is the one it must not do twice. +for n, n_spikes, duration in ((30, 40, 120.0), (60, 8, 300.0), (12, 1, 60.0), (31, 60, 90.0)): r = np.random.default_rng(n) trains = {i: np.sort(r.uniform(0, duration, r.integers(0, n_spikes + 1))) for i in range(n)} # A couple of empty trains, which is the branch that returns NaN. trains[0] = np.array([]) trains[n - 1] = np.array([]) + # CAT-NAP peak times can arrive unsorted, and run_P is order-dependent. + trains[2] = r.permutation(r.uniform(0, duration, 25)) for lag_ms in (10.0, 1000.0): - got = get_sttc(trains, n, lag_ms, duration) - want = np.full((n, n), np.nan) - for i in range(n): - for j in range(i + 1, n): - c = sttc_pair(trains[i], trains[j], lag_ms / 1000.0, 0.0, duration) - want[i, j] = c - want[j, i] = c - want[want < 0] = 0.0 - want[np.isnan(want)] = 0.0 - check(f"get_sttc n={n} lag={lag_ms:.0f}ms", np.array_equal(got, want), - f"max diff {np.max(np.abs(got - want)) if got.shape == want.shape else 'shape'}") + want = sttc_reference(trains, n, lag_ms, duration) + for threads in (1, numba.config.NUMBA_NUM_THREADS): + numba.set_num_threads(threads) + got = get_sttc(trains, n, lag_ms, duration) + check(f"get_sttc n={n} lag={lag_ms:.0f}ms threads={threads}", + np.array_equal(got, want), + f"max diff {np.max(np.abs(got - want)) if got.shape == want.shape else 'shape'}") +numba.set_num_threads(numba.config.NUMBA_NUM_THREADS) +check("get_sttc n=0", get_sttc({}, 0, 10.0, 60.0).shape == (0, 0)) + + +print() +print("[4] randmio_und_v2 / latmio_und_v2 — compiled vs the generator version") + + +def _ref_valid_quad_stream(rng, k, i_arr, j_arr, batch_size=8192): + """The pre-numba ``_valid_quad_stream``, verbatim.""" + budget = min(2_000_000, max(10_000, 200 * k * k)) + rejected = 0 + while True: + e1_batch = rng.integers(0, k, size=batch_size) + e2_batch = rng.integers(0, k, size=batch_size) + for e1, e2 in zip(e1_batch.tolist(), e2_batch.tolist()): + if e1 != e2: + a, b = i_arr[e1], j_arr[e1] + c, d = i_arr[e2], j_arr[e2] + if a != c and a != d and b != c and b != d: + rejected = 0 + yield e1, e2, a, b, c, d + continue + rejected += 1 + if rejected >= budget: + yield None + rejected = 0 + + +def _ref_coin_flip_stream(rng, batch_size=8192): + while True: + for v in rng.random(batch_size).tolist(): + yield v > 0.5 + + +def ref_randmio_und_v2(w, iterations, rng): + """The pre-numba ``randmio_und_v2``, verbatim.""" + r = w.astype(float).copy() + n = r.shape[0] + i_arr_np, j_arr_np = np.nonzero(np.tril(r)) + i_arr, j_arr = i_arr_np.tolist(), j_arr_np.tolist() + k = len(i_arr) + if k < 2: + return r + max_attempts = round(n * k / (n * (n - 1))) + quads = _ref_valid_quad_stream(rng, k, i_arr, j_arr) + flips = _ref_coin_flip_stream(rng) + for _ in range(iterations): + for _attempt in range(max_attempts + 1): + quad = next(quads) + if quad is None: + return r + e1, e2, a, b, c, d = quad + if next(flips): + i_arr[e2], j_arr[e2] = d, c + c, d = d, c + if not (r[a, d] or r[c, b]): + r[a, d], r[a, b] = r[a, b], 0.0 + r[d, a], r[b, a] = r[b, a], 0.0 + r[c, b], r[c, d] = r[c, d], 0.0 + r[b, c], r[d, c] = r[d, c], 0.0 + j_arr[e1] = d + j_arr[e2] = b + break + return r + + +def ref_latmio_und_v2(w, iterations, d, rng): + """The pre-numba ``latmio_und_v2``, verbatim.""" + n = w.shape[0] + ind_rp = rng.permutation(n) + r = w[np.ix_(ind_rp, ind_rp)].astype(float).copy() + i_arr_np, j_arr_np = np.nonzero(np.tril(r)) + i_arr, j_arr = i_arr_np.tolist(), j_arr_np.tolist() + k = len(i_arr) + if k >= 2: + max_attempts = round(n * k / (n * (n - 1) / 2)) + quads = _ref_valid_quad_stream(rng, k, i_arr, j_arr) + flips = _ref_coin_flip_stream(rng) + for _ in range(iterations): + exhausted = False + for _attempt in range(max_attempts + 1): + quad = next(quads) + if quad is None: + exhausted = True + break + e1, e2, a, b, c, dd = quad + if next(flips): + i_arr[e2], j_arr[e2] = dd, c + c, dd = dd, c + if not (r[a, dd] or r[c, b]): + if (d[a, b] * r[a, b] + d[c, dd] * r[c, dd]) >= ( + d[a, dd] * r[a, b] + d[c, b] * r[c, dd] + ): + r[a, dd], r[a, b] = r[a, b], 0.0 + r[dd, a], r[b, a] = r[b, a], 0.0 + r[c, b], r[c, dd] = r[c, dd], 0.0 + r[b, c], r[dd, c] = r[dd, c], 0.0 + j_arr[e1] = dd + j_arr[e2] = b + break + if exhausted: + break + ind_rp_reverse = np.argsort(ind_rp) + return r[np.ix_(ind_rp_reverse, ind_rp_reverse)] + + +def profile_distance(w: np.ndarray) -> np.ndarray: + """``squareform(pdist(adjM))``, what step 4 passes latmio as ``d``.""" + diff = w[:, None, :] - w[None, :, :] + return np.sqrt(np.sum(diff * diff, axis=-1)) + + +star = np.zeros((6, 6)) +star[0, 1:] = star[1:, 0] = 1.0 +with_nan = random_graph(20, 0.4, 7) +with_nan[2, 5] = with_nan[5, 2] = np.nan +one_edge = np.zeros((10, 10)) +one_edge[1, 4] = one_edge[4, 1] = 0.5 +swap_cases = [(f"n={n} density={dens}", random_graph(n, dens, seed)) + for n, dens, seed in ((5, 0.9, 1), (12, 0.3, 2), (40, 0.05, 3), + (64, 0.1, 4), (64, 1.0, 5), (150, 0.06, 6))] +swap_cases += [("triangle (no valid swap)", np.ones((3, 3)) - np.eye(3)), + ("star (no valid swap)", star), + ("a NaN weight", with_nan), + ("one edge", one_edge)] +for label, w in swap_cases: + d = np.nan_to_num(profile_distance(np.nan_to_num(w))) + for name, fast, slow, args in ( + ("randmio_und_v2", randmio_und_v2, ref_randmio_und_v2, (w, 200)), + ("latmio_und_v2", latmio_und_v2, ref_latmio_und_v2, (w, 400, d)), + ): + r_fast = np.random.default_rng(11) + r_slow = np.random.default_rng(11) + got = fast(*args, rng=r_fast) + want = slow(*args, rng=r_slow) + check(f"{name}: {label}", + np.array_equal(got, want, equal_nan=True) + and r_fast.bit_generator.state == r_slow.bit_generator.state, + "networks or generator state differ") + + +print() +print("[5] cal_nmf — threaded fits vs the serial loops") + + +def ref_cal_nmf_curves(trains, counts, duration, rng): + """The pre-threading search and sweep loops, on the unchanged ``_nnmf``.""" + n_bins = max(int(round(10.0 * duration)), 1) + x = nmf._bin_spike_times(trains, duration, n_bins) + rand = nmf._bin_spike_times([nmf.randomise_spike_train(t, duration, rng) for t in trains], + duration, n_bins) + residual, rand_residual, k, per_component = 0.0, 1.0, 1, [] + while residual < rand_residual and k <= x.shape[1]: + _, _, residual = nmf._nnmf(x, k, rng) + _, _, rand_residual = nmf._nnmf(rand, k, rng) + per_component.append(rand_residual) + k += 1 + active = counts > 1 + xa = x[:, active] + res, ve = [], [] + for kk in range(1, int(active.sum()) + 1): + w, h, d = nmf._nnmf(xa, kk, rng) + res.append(d) + ve.append(nmf._variance_explained(xa, w, h)) + return k - 1, np.array(per_component), np.array(res), np.array(ve) + + +def structured_trains(n: int, duration: float, seed: int) -> list[np.ndarray]: + """Units sharing a few drivers, so the component search goes past rank 1.""" + r = np.random.default_rng(seed) + drivers = [np.sort(r.uniform(0, duration, 300)) for _ in range(3)] + out = [] + for ch in range(n): + if r.random() < 0.15: + out.append(np.array([])) + continue + own = r.uniform(0, duration, int(r.integers(5, 400))) + shared = drivers[ch % 3] + r.normal(0, 0.01, 300) + out.append(np.sort(np.clip(np.concatenate([own, shared[r.random(300) < 0.5]]), + 0, duration - 1e-6))) + return out + + +for seed, n in ((0, 20), (3, 13), (5, 18), (1, 1)): + trains = structured_trains(n, 120.0, seed) + counts = np.array([len(t) for t in trains]) + r_ref = np.random.default_rng(seed) + k_ref, per_ref, res_ref, ve_ref = ref_cal_nmf_curves(trains, counts, 120.0, r_ref) + for workers in (1, 3, 8): + r_new = np.random.default_rng(seed) + got = nmf.cal_nmf(trains, counts, 120.0, 10.0, 25000.0, rng=r_new, max_workers=workers) + same = (got["num_nnmf_components"] == k_ref + and np.array_equal(got["randResidualPerComponent"], per_ref) + and np.array_equal(got["nnmf_residuals"], res_ref) + and np.array_equal(got["nnmf_var_explained"], ve_ref) + and r_new.bit_generator.state == r_ref.bit_generator.state) + check(f"cal_nmf n={n} workers={workers} (k={k_ref})", same, + "curves, component count or generator state differ") + +# The early stop must fit a prefix of the full sweep, rank for rank, then stop +# where _levelled_off first says so, leaving the rest NaN. +trains = structured_trains(60, 120.0, 9) +counts = np.array([len(t) for t in trains]) +full = nmf.cal_nmf(trains, counts, 120.0, 10.0, 25000.0, rng=np.random.default_rng(2), max_workers=4) +first_95 = int(np.argmax(full["nnmf_var_explained"] > 0.95)) + 1 +stops = [] +for gain in (0.002, 0.05, 0.3): + part = nmf.cal_nmf(trains, counts, 120.0, 10.0, 25000.0, rng=np.random.default_rng(2), + max_workers=4, plateau_gain=gain) + ve = part["nnmf_var_explained"] + fitted = int(np.sum(~np.isnan(ve))) + expected = len(ve) + for r in range(1, len(ve) + 1): + reached = bool(np.any(full["nnmf_var_explained"][:r] > 0.95)) + if nmf._levelled_off(full["nnmf_var_explained"][:r], gain, reached): + expected = r + break + check(f"plateau_gain={gain}: stops at rank {expected} of {len(ve)}", + fitted == expected + and np.array_equal(ve[:fitted], full["nnmf_var_explained"][:fitted]) + and np.all(np.isnan(ve[fitted:])) + and part["num_nnmf_components"] == full["num_nnmf_components"], + f"fitted {fitted}, expected {expected}") + stops.append(fitted) +# The large gain has to stop on a plateau, before the curve reaches 95%, or the +# plateau branch of _levelled_off has gone untested. +check(f"a plateau stop comes before the 95% crossing (rank {first_95})", + min(stops) < first_95, f"stops {stops}") + + +print() +print("[6] null_model_und_sign — threaded batches vs drawing as it goes") + + +def signed_graph(n: int, density: float, seed: int) -> np.ndarray: + r = np.random.default_rng(seed) + w = r.random((n, n)) * (r.random((n, n)) < density) * np.where(r.random((n, n)) < 0.3, -1, 1) + w = np.triu(w, 1) + return w + w.T + + +null_cases = [("n=3 (too small to rewire)", random_graph(3, 1.0, 0)), + ("complete n=20 (not rewired)", random_graph(20, 1.0, 1)), + ("n=30 signed", signed_graph(30, 0.2, 2)), + ("n=64 density 0.1", random_graph(64, 0.1, 3)), + ("n=150 density 0.06", random_graph(150, 0.06, 4))] +for label, w in null_cases: + r_ref = np.random.default_rng(5) + want = [_null_from_plan(w, 5, 0.1, None, r_ref) for _ in range(10)] + for workers in (1, 6): + r_new = np.random.default_rng(5) + got = map_null_models_und_sign(w, 10, lambda x: x, rng=r_new, workers=workers) + check(f"null models {label}, workers={workers}", + all(np.array_equal(a, b) for a, b in zip(want, got)) + and r_new.bit_generator.state == r_ref.bit_generator.state, + "networks or generator state differ") + +w = random_graph(80, 0.1, 6) +ci = np.random.default_rng(6).integers(1, 5, 80) +_threads = nm.compute_threads +try: + nm.compute_threads = lambda: 1 + serial = nm.participation_coef_norm(w, ci, 40, rng=np.random.default_rng(7)) + nm.compute_threads = lambda: 8 + threaded = nm.participation_coef_norm(w, ci, 40, rng=np.random.default_rng(7)) +finally: + nm.compute_threads = _threads +check("participation_coef_norm: 1 thread vs 8", + all(np.array_equal(a, b) for a, b in zip(serial, threaded))) print() diff --git a/python/test_step4_optional_metrics.py b/python/test_step4_optional_metrics.py index 8fde7f20..6d9612be 100644 --- a/python/test_step4_optional_metrics.py +++ b/python/test_step4_optional_metrics.py @@ -258,17 +258,29 @@ def _gui_checks() -> list[Check]: checks.append(("and writes back what was ticked", out.compute_nmf and not out.compute_eff_rank, "")) - # CAT-NAP reads its own NMF switch, so these must not look live there. + # CAT-NAP reads its own NMF switch, so these must not look live there. The + # NMF sweep setting is read by both pipelines, so it must stay live. panel.set_pipeline("catnap") - checks.append(("in CAT-NAP mode the box is disabled, not silently ignored", - not panel._metrics_box.isEnabled() + checks.append(("in CAT-NAP mode the switches are disabled, not silently ignored", + not panel.compute_nmf.isEnabled() + and not panel.compute_eff_rank.isEnabled() and "CAT-NAP tab" in panel._metrics_box.title(), panel._metrics_box.title())) + checks.append(("but the NMF sweep setting, which CAT-NAP reads, stays live", + panel.nmf_plateau_gain.isEnabled(), "")) panel.set_pipeline("meanap") checks.append(("and live again for an ephys run", - panel._metrics_box.isEnabled() + panel.compute_nmf.isEnabled() and panel.compute_eff_rank.isEnabled() and panel._metrics_box.title() == "Network metrics", panel._metrics_box.title())) + + # The sweep setting round-trips, including 0 (the "off" special value). + for gain in (0.0, 0.02): + panel.load(Params(nmf_plateau_gain=gain)) + out = Params() + panel.save(out) + checks.append((f"nmf_plateau_gain={gain} survives load and save", + out.nmf_plateau_gain == gain, str(out.nmf_plateau_gain))) return checks diff --git a/src/meanap/app_smoke.py b/src/meanap/app_smoke.py index 658816d1..ce2d9cac 100644 --- a/src/meanap/app_smoke.py +++ b/src/meanap/app_smoke.py @@ -81,7 +81,16 @@ def check_numba() -> str: adj = np.random.default_rng(0).random((20, 20)) adj = np.triu(adj, 1) + np.triu(adj, 1).T null_models.randmio_und_signed(adj, 2, rng=np.random.default_rng(1)) - return f"numba {numba.__version__}" + + # STTC's sweep is the one ``parallel=True`` kernel: it needs a threading + # layer, which numba loads by name and a frozen build can leave out. + from meanap.pipeline import sttc + + trains = {i: np.sort(np.random.default_rng(i).uniform(0, 10, 50)) for i in range(8)} + sttc.get_sttc(trains, 8, 10.0, 10.0) + if sttc._PARALLEL_BROKEN: + raise RuntimeError("numba has no threading layer; STTC fell back to one core") + return f"numba {numba.__version__}, threading layer {numba.threading_layer()}" def check_pipeline() -> str: diff --git a/src/meanap/catnap/pipeline.py b/src/meanap/catnap/pipeline.py index 46be2de8..0ef4e7c7 100644 --- a/src/meanap/catnap/pipeline.py +++ b/src/meanap/catnap/pipeline.py @@ -1033,7 +1033,8 @@ def _lag_independent_metrics( out.update(cal_nmf( spike_times, np.asarray(_spike_counts(res, params.twop_activity)), duration_s, params.nmf_downsample_freq, res.fs, - include_nmf_components=params.include_nmf_components, rng=rng)) + include_nmf_components=params.include_nmf_components, rng=rng, + plateau_gain=params.nmf_plateau_gain)) except Exception as e: log(f" [{name}] WARNING: could not compute NMF components: {e}") diff --git a/src/meanap/gui/panels/connectivity.py b/src/meanap/gui/panels/connectivity.py index fa50c469..68f7831d 100644 --- a/src/meanap/gui/panels/connectivity.py +++ b/src/meanap/gui/panels/connectivity.py @@ -163,8 +163,30 @@ def __init__(self, parent: QWidget | None = None) -> None: "NMF, but the same kind of measure, so it can be left out " "the same way.") + self.nmf_plateau_gain = QDoubleSpinBox() + self.nmf_plateau_gain.setRange(0.0, 1.0) + self.nmf_plateau_gain.setDecimals(3) + self.nmf_plateau_gain.setSingleStep(0.001) + self.nmf_plateau_gain.setSpecialValueText("off (fit every rank)") + self.nmf_plateau_gain.setValue(0.005) + set_tooltip(self.nmf_plateau_gain, + "Where NMF's variance-explained curve stops.\n\n" + "The curve fits one factorisation per rank, up to the " + "number of active nodes, and the highest ranks cost the " + "most: with hundreds of sorted units they are most of " + "step 4. The sweep stops once variance explained passes " + "95%, or rises by less than this over 10 ranks. Ranks it " + "skips are left empty in the curve; num_nnmf_components " + "and nComponentsRelNS do not depend on it.\n\n" + "0 fits every rank, as MATLAB does. Applies to CAT-NAP's " + "NMF too.") + form3.addRow("NMF dimensionality", self.compute_nmf) form3.addRow("Effective rank", self.compute_eff_rank) + metrics_advanced = AdvancedSection() + metrics_advanced.form().addRow("Stop NMF sweep below gain", + self.nmf_plateau_gain) + form3.addRow(metrics_advanced) self._thr_box = thr_box self._metrics_box = metrics_box @@ -262,14 +284,18 @@ def set_pipeline(self, mode_key: str) -> None: "Measured in spikes/s over the whole recording. 0 keeps " "every electrode that has at least one edge, which is the " "default; MEApipeline.m's own example runs use 0.01.") - self._metrics_box.setEnabled(not catnap) + # The two switches are ephys-only; the NMF sweep setting is read by + # both pipelines, so it stays live. + self.compute_nmf.setEnabled(not catnap) + self.compute_eff_rank.setEnabled(not catnap) self._metrics_box.setTitle( - "Network metrics · set on the CAT-NAP tab" if catnap + "Network metrics · NMF is switched on the CAT-NAP tab" if catnap else "Network metrics") set_tooltip(self._metrics_box, - "A CAT-NAP run takes its NMF setting from the CAT-NAP " - "tab, and always computes effective rank, so nothing here " - "is read." if catnap else "") + "A CAT-NAP run takes its NMF switch from the CAT-NAP tab " + "and always computes effective rank, so the two switches " + "here are not read. The NMF sweep setting under Advanced " + "is." if catnap else "") def retune_lags_for_mode(self, from_mode: str, to_mode: str) -> tuple[int, ...] | None: """Move the lags onto *to_mode*'s defaults, unless they were edited. @@ -315,6 +341,7 @@ def load(self, params: Params) -> None: params.exclude_edges_below_threshold) self.compute_nmf.setChecked(params.compute_nmf) self.compute_eff_rank.setChecked(params.compute_eff_rank) + self.nmf_plateau_gain.setValue(params.nmf_plateau_gain) def save(self, params: Params) -> None: params.func_con_lag_val = self._parse_lags() @@ -329,3 +356,4 @@ def save(self, params: Params) -> None: self.exclude_edges_below_threshold.isChecked()) params.compute_nmf = self.compute_nmf.isChecked() params.compute_eff_rank = self.compute_eff_rank.isChecked() + params.nmf_plateau_gain = self.nmf_plateau_gain.value() diff --git a/src/meanap/params.py b/src/meanap/params.py index c6e5199d..a04ffd38 100644 --- a/src/meanap/params.py +++ b/src/meanap/params.py @@ -237,6 +237,14 @@ class Params: eff_rank_downsample_freq: float = 10.0 nmf_downsample_freq: float = 10.0 include_nmf_components: bool = False + # NMF's per-rank sweep (the variance-explained curve) fits every rank up to + # the number of active nodes in calNMF.m, and the highest ranks cost the + # most. With this above 0 the sweep stops once variance explained passes + # 95% or rises by less than this over 10 ranks; the ranks it skips are NaN + # in nnmf_residuals / nnmf_var_explained. num_nnmf_components and + # nComponentsRelNS do not depend on it. 0 fits every rank, as MATLAB does. + # Read by both pipelines. + nmf_plateau_gain: float = 0.005 # ── Plotting ───────────────────────────────────────────────────────────── fig_ext: list[str] = field(default_factory=lambda: [".png"]) diff --git a/src/meanap/pipeline/network_metrics.py b/src/meanap/pipeline/network_metrics.py index e00f8fbd..72747e43 100644 --- a/src/meanap/pipeline/network_metrics.py +++ b/src/meanap/pipeline/network_metrics.py @@ -55,7 +55,8 @@ from scipy.linalg import schur from scipy.sparse.linalg import svds -from meanap.pipeline.null_models import null_model_und_sign +from meanap.pipeline.null_models import map_null_models_und_sign +from meanap.pipeline.parallel import compute_threads # Same optional-numba pattern as ``null_models.py`` / ``sttc.py``: numba is a # declared dependency, but the import is guarded so a build without it still @@ -426,16 +427,22 @@ def participation_coef_norm( within_mod_k[mask] = w[np.ix_(mask, mask)].sum(axis=1) between_mod_k = ko - within_mod_k - kc2_rnd = np.zeros((n, n_iter)) - for it in range(n_iter): - w_rnd = null_model_und_sign(w, bin_swaps=5, rng=rng) + def kc2_column(w_rnd: np.ndarray) -> np.ndarray: gc_rnd = (w_rnd != 0) @ np.diag(ci) kc2_rnd_loop = np.zeros(n) with np.errstate(divide="ignore", invalid="ignore"): for i in range(1, ci.max() + 1): term = (w * (gc == i)).sum(axis=1) / ko - (w_rnd * (gc_rnd == i)).sum(axis=1) / ko kc2_rnd_loop += term**2 - kc2_rnd[:, it] = np.sqrt(0.5 * kc2_rnd_loop) + return np.sqrt(0.5 * kc2_rnd_loop) + + # The null models are the cost here (each rewiring grows as n^3), and + # they are independent, so they are built on threads; the result and the + # generator state afterwards are those of a sequential loop (see + # map_null_models_und_sign). + columns = map_null_models_und_sign(w, n_iter, kc2_column, bin_swaps=5, rng=rng, + workers=compute_threads()) + kc2_rnd = np.column_stack(columns) if columns else np.zeros((n, 0)) with np.errstate(invalid="ignore"): pc_norm = 1.0 - np.median(kc2_rnd, axis=1) diff --git a/src/meanap/pipeline/nmf.py b/src/meanap/pipeline/nmf.py index f87db93a..f073706f 100644 --- a/src/meanap/pipeline/nmf.py +++ b/src/meanap/pipeline/nmf.py @@ -30,10 +30,18 @@ from __future__ import annotations import warnings +from collections import deque +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack +from itertools import islice import numpy as np from sklearn.decomposition import NMF from sklearn.exceptions import ConvergenceWarning +from threadpoolctl import threadpool_limits + +from meanap.pipeline.parallel import compute_threads def randomise_spike_train( @@ -68,7 +76,17 @@ def _bin_spike_times(spike_times_list: list[np.ndarray], duration_s: float, n_bi return out -def _nnmf(x: np.ndarray, k: int, rng: np.random.Generator) -> tuple[np.ndarray, np.ndarray, float]: +def _draw_seed(rng: np.random.Generator) -> int: + """One fit's ``random_state``, drawn exactly as ``_nnmf`` always drew it. + + Seeds are drawn one call at a time and in rank order, never as an array: + an array draw takes a different path through the generator, so the seeds + (and every fit) would change. + """ + return int(rng.integers(0, 2**31 - 1)) + + +def _nnmf_seeded(x: np.ndarray, k: int, seed: int) -> tuple[np.ndarray, np.ndarray, float]: """Rank-``k`` NNMF of ``x``, matching MATLAB's ``[W,H,D] = nnmf(A,k)`` output convention: ``D`` is the root-mean-square residual ``norm(A - W*H, 'fro') / sqrt(m*n)``, not sklearn's raw Frobenius error. @@ -87,16 +105,9 @@ def _nnmf(x: np.ndarray, k: int, rng: np.random.Generator) -> tuple[np.ndarray, """ m, n = x.shape k = max(1, min(k, m, n)) - seed = int(rng.integers(0, 2**31 - 1)) try: model = NMF(n_components=k, init="random", solver="cd", max_iter=100, random_state=seed) - with warnings.catch_warnings(): - # Not reaching full convergence within a bounded iteration count - # is expected/tolerated here, same as MATLAB's own ALS default - # — this metric is inherently approximate, not exactly - # reproducible even between two MATLAB runs on the same data. - warnings.simplefilter("ignore", category=ConvergenceWarning) - w = model.fit_transform(x) + w = model.fit_transform(x) h = model.components_ d = model.reconstruction_err_ / np.sqrt(m * n) except Exception: @@ -106,6 +117,116 @@ def _nnmf(x: np.ndarray, k: int, rng: np.random.Generator) -> tuple[np.ndarray, return w, h, float(d) +def _nnmf(x: np.ndarray, k: int, rng: np.random.Generator) -> tuple[np.ndarray, np.ndarray, float]: + """:func:`_nnmf_seeded` with the next seed from ``rng``.""" + return _fit_all([(x, k, _draw_seed(rng))], 1)[0] + + +def _quiet_single_blas(): + """Context for running fits: BLAS on one thread, ConvergenceWarning off. + + BLAS is held to one thread for every fit, threaded or not. A multi-threaded + BLAS splits its sums differently for different thread counts, which moves + the factors in the last bits, and the fits amplify that. With one BLAS + thread a fit no longer depends on how many cores the machine has or on how + many fits run at once. + + Not reaching full convergence within the bounded iteration count is + expected and tolerated, as with MATLAB's own ALS default, so sklearn's + ConvergenceWarning is silenced. That has to happen in the calling thread, + once: ``warnings.catch_warnings`` swaps process-wide state, so entering it + from several fit threads at once can leak warnings or leave the filter + changed. + """ + stack = ExitStack() + stack.enter_context(warnings.catch_warnings()) + stack.enter_context(threadpool_limits(limits=1, user_api="blas")) + warnings.simplefilter("ignore", category=ConvergenceWarning) + return stack + + +def _fit_all(jobs: list[tuple[np.ndarray, int, int]], workers: int) -> list: + """Run ``_nnmf_seeded`` on each ``(x, k, seed)``; results in job order. + + Every fit has its own seed, so running them on threads changes nothing but + the wall time. sklearn's coordinate-descent update releases the GIL, so the + threads do run in parallel. + """ + with _quiet_single_blas(): + if workers <= 1 or len(jobs) <= 1: + return [_nnmf_seeded(*job) for job in jobs] + with ThreadPoolExecutor(max_workers=min(workers, len(jobs))) as ex: + return list(ex.map(lambda job: _nnmf_seeded(*job), jobs)) + + +def _sweep_fit(x: np.ndarray, k: int, seed: int) -> tuple[int, np.ndarray, np.ndarray, float, float]: + """One rank of the per-rank sweep: the rank, its fit and variance explained.""" + w, h, res = _nnmf_seeded(x, k, seed) + return k, w, h, res, _variance_explained(x, w, h) + + +def _fit_in_order(jobs: list[tuple[np.ndarray, int, int]], workers: int, + consume: Callable[[tuple], bool]) -> None: + """Run ``_sweep_fit`` on each job, handing results to ``consume`` in order. + + Up to ``workers`` fits run at once, and a new one starts as soon as the + oldest is consumed, so a slow fit does not hold up a whole batch. When + ``consume`` returns False the remaining jobs are dropped (fits already + running finish and are discarded). + """ + with _quiet_single_blas(): + if workers <= 1: + for job in jobs: + if not consume(_sweep_fit(*job)): + return + return + with ThreadPoolExecutor(max_workers=workers) as ex: + queued = iter(jobs) + running = deque(ex.submit(_sweep_fit, *job) for job in islice(queued, workers)) + try: + while running: + result = running.popleft().result() + nxt = next(queued, None) + if nxt is not None: + running.append(ex.submit(_sweep_fit, *nxt)) + if not consume(result): + return + finally: + for fut in running: + fut.cancel() + + +def _variance_explained(x: np.ndarray, w: np.ndarray, h: np.ndarray) -> float: + """``calNMF.m``'s variance explained: 1 - SS_res / SS_tot about the grand mean.""" + ss_res = np.sum((w @ h - x) ** 2) + ss_tot = np.sum((x - x.mean()) ** 2) + return (1.0 - ss_res / ss_tot) if ss_tot > 0 else 0.0 + + +#: Ranks over which ``cal_nmf``'s early stop judges that variance explained +#: has levelled off. +PLATEAU_WINDOW = 10 + + +def _levelled_off(var_explained: np.ndarray, gain: float, threshold_reached: bool) -> bool: + """Whether the per-rank sweep can stop after the ranks fitted so far. + + True once the curve has passed the 95% threshold, or once the best value + over the last ``PLATEAU_WINDOW`` ranks beats the best before them by less + than ``gain``. The maxima, not the last values, are compared because the + curve is not monotonic: each rank is an independent fit from a random + start, and a high rank that fails to converge can drop well below a lower + one. + """ + if threshold_reached: + return True + r = var_explained.shape[0] + if r <= PLATEAU_WINDOW: + return False + return (np.max(var_explained[r - PLATEAU_WINDOW:]) + - np.max(var_explained[:r - PLATEAU_WINDOW])) < gain + + def cal_nmf( spike_times_list: list[np.ndarray], spike_counts: np.ndarray, @@ -115,6 +236,8 @@ def cal_nmf( min_spike_count: int = 1, include_nmf_components: bool = False, rng: np.random.Generator | None = None, + plateau_gain: float = 0.0, + max_workers: int | None = None, ) -> dict: """Non-negative matrix factorization dimensionality metrics for one recording (lag-independent — call once per recording, not once per lag, @@ -127,9 +250,19 @@ def cal_nmf( matching MATLAB's ``Params.includeNMFcomponents``), ``nmfFactors``, ``nmfWeights``, ``downSampleSpikeMatrix``, ``nmfFactorsVarThreshold``, ``nmfWeightsVarThreshold``. + + ``plateau_gain`` > 0 ends the per-rank sweep early, once variance + explained has risen by less than that over the last + ``PLATEAU_WINDOW`` ranks (or has passed 95%); the ranks not fitted are + NaN in ``nnmf_residuals``/``nnmf_var_explained``, and if 95% was never + reached the ``*VarThreshold`` factors come from the last rank fitted + rather than the last rank. 0 fits every rank, as ``calNMF.m`` does. ``max_workers`` caps the threads the fits run on + (default: :func:`~meanap.pipeline.parallel.compute_threads`); the result + is the same for any value. """ if rng is None: rng = np.random.default_rng() + workers = compute_threads() if max_workers is None else max(1, max_workers) downsample_freq = min(downsample_freq, fs) n_bins = max(int(round(downsample_freq * duration_s)), 1) @@ -146,15 +279,39 @@ def cal_nmf( # rather than fitting noise as readily as they'd fit shuffled data. ── nmf_factors: np.ndarray | None = None nmf_weights: np.ndarray | None = None - residual = 0.0 - rand_residual = 1.0 - k = 1 rand_residual_per_component: list[float] = [] - while residual < rand_residual and k <= n_channels: - nmf_factors, nmf_weights, residual = _nnmf(down_sample_spike_matrix, k, rng) - _, _, rand_residual = _nnmf(rand_spike_matrix, k, rng) - rand_residual_per_component.append(rand_residual) - k += 1 + k = 1 + block = 1 + searching = n_channels >= 1 + while searching: + # Fit ranks k..k+block-1 at once, real and shuffled side by side. The + # loop stops at the first rank whose residual no longer beats the + # shuffled one, so fits past it are wasted: the block starts at one + # rank and doubles while the search continues, up to the thread + # count, so a search that ends at rank 1 or 2 costs what it did serially. + ranks = list(range(k, min(k + block, n_channels + 1))) + state = rng.bit_generator.state + seeds = [(_draw_seed(rng), _draw_seed(rng)) for _ in ranks] + fits = _fit_all([job for kk, (s_real, s_rand) in zip(ranks, seeds) + for job in ((down_sample_spike_matrix, kk, s_real), + (rand_spike_matrix, kk, s_rand))], workers) + used = 0 + for idx, kk in enumerate(ranks): + nmf_factors, nmf_weights, residual = fits[2 * idx] + rand_residual = fits[2 * idx + 1][2] + rand_residual_per_component.append(rand_residual) + used += 1 + k = kk + 1 + if not (residual < rand_residual and k <= n_channels): + searching = False + break + if used < len(ranks): + # Put the generator back as if only the ranks used had drawn + # seeds, so everything after sees the stream the serial loop left. + rng.bit_generator.state = state + for _ in range(2 * used): + _draw_seed(rng) + block = min(2 * block, max(1, workers)) num_nnmf_components = k - 1 active_electrodes = spike_counts > min_spike_count @@ -169,28 +326,35 @@ def cal_nmf( # ── Per-rank sweep over active electrodes only, for the "how many # components until we explain 95% of variance" curve. ── down_sample_active = down_sample_spike_matrix[:, active_electrodes] - nnmf_residuals = np.zeros(network_size) - nnmf_var_explained = np.zeros(network_size) + # NaN until fitted: an early stop leaves the ranks it skipped as NaN. + nnmf_residuals = np.full(network_size, np.nan) + nnmf_var_explained = np.full(network_size, np.nan) var_explained_threshold = 0.95 threshold_reached = False nmf_factors_var_threshold: np.ndarray | None = None nmf_weights_var_threshold: np.ndarray | None = None - for kk in range(1, network_size + 1): - w, h, res = _nnmf(down_sample_active, kk, rng) - nnmf_residuals[kk - 1] = res - - predicted = w @ h - ss_res = np.sum((predicted - down_sample_active) ** 2) - grand_mean = down_sample_active.mean() - ss_tot = np.sum((down_sample_active - grand_mean) ** 2) - var_explained = (1.0 - ss_res / ss_tot) if ss_tot > 0 else 0.0 - nnmf_var_explained[kk - 1] = var_explained + # Every rank's seed is drawn up front, in rank order, so a rank's fit is the + # same however many threads run the sweep and wherever it stops. + sweep_seeds = [_draw_seed(rng) for _ in range(network_size)] + last: list = [None, None] + def take(fit: tuple) -> bool: + nonlocal threshold_reached, nmf_factors_var_threshold, nmf_weights_var_threshold + rank, w_r, h_r, res, var_explained = fit + last[:] = [w_r, h_r] + nnmf_residuals[rank - 1] = res + nnmf_var_explained[rank - 1] = var_explained if var_explained > var_explained_threshold and not threshold_reached: threshold_reached = True - nmf_factors_var_threshold = w - nmf_weights_var_threshold = h + nmf_factors_var_threshold = w_r + nmf_weights_var_threshold = h_r + return not (plateau_gain > 0 and _levelled_off( + nnmf_var_explained[:rank], plateau_gain, threshold_reached)) + + _fit_in_order([(down_sample_active, r, sweep_seeds[r - 1]) for r in range(1, network_size + 1)], + workers, take) + w, h = last if network_size > 0 and not threshold_reached: nmf_factors_var_threshold = w diff --git a/src/meanap/pipeline/null_models.py b/src/meanap/pipeline/null_models.py index 3f08b701..5c092ccb 100644 --- a/src/meanap/pipeline/null_models.py +++ b/src/meanap/pipeline/null_models.py @@ -16,10 +16,15 @@ from __future__ import annotations -from typing import Iterator +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from typing import TypeVar import numpy as np +T = TypeVar("T") + # ``randmio_und_signed`` is the pipeline's single hottest Step-4 function # (profiled at ~46% of Step 4 — 159M inner iterations of a sequential swap # loop that can't be vectorized, since each accepted swap mutates the matrix @@ -37,12 +42,16 @@ if _HAVE_NUMBA: - @njit(cache=True) + @njit(cache=True, nogil=True) def _randmio_signed_core(r, sgn, total_iter, max_attempts, seed): # pragma: no cover - jitted """In-place degree-preserving signed swaps on ``r`` (and its sign matrix ``sgn``). Compiled; uses numba's own RNG (seeded per call). ``sgn`` is kept consistent with ``r`` so sign comparisons are integer - lookups rather than repeated ``np.sign`` calls.""" + lookups rather than repeated ``np.sign`` calls. + + Releases the GIL, so several null models can rewire on threads at + once. numba keeps one RNG state per thread and this seeds it before + drawing, so a call's result does not depend on what else is running.""" np.random.seed(seed) n = r.shape[0] for _ in range(total_iter): @@ -74,6 +83,17 @@ def _randmio_signed_core(r, sgn, total_iter, max_attempts, seed): # pragma: no break +def _randmio_signed_seeded(r: np.ndarray, iterations: int, seed: int) -> np.ndarray: + """The numba path of :func:`randmio_und_signed`, on a float copy ``r`` + (rewired in place and returned) with its seed already drawn.""" + n = r.shape[0] + total_iter = int(iterations * n * (n - 1) / 2) + max_attempts = round(n / 2) + sgn = np.sign(r).astype(np.int64) + _randmio_signed_core(r, sgn, total_iter, max_attempts, seed) + return r + + def randmio_und_signed( w: np.ndarray, iterations: int, rng: np.random.Generator | None = None, ) -> np.ndarray: @@ -102,10 +122,7 @@ def randmio_und_signed( max_attempts = round(n / 2) if _HAVE_NUMBA: - sgn = np.sign(r).astype(np.int64) - seed = int(rng.integers(0, 2**31 - 1)) - _randmio_signed_core(r, sgn, total_iter, max_attempts, seed) - return r + return _randmio_signed_seeded(r, iterations, int(rng.integers(0, 2**31 - 1))) # ── Pure-NumPy fallback ────────────────────────────────────────────────── # rng.choice(n, size=4, replace=False) has surprisingly high per-call @@ -183,23 +200,58 @@ def _stable_ranks(values: np.ndarray, ranks: np.ndarray) -> np.ndarray: return picked -def null_model_und_sign( - w: np.ndarray, - bin_swaps: int = 5, - wei_freq: float = 0.1, - rng: np.random.Generator | None = None, -) -> np.ndarray: - """Randomize an undirected network, preserving degree and (approximately) - strength distributions. +@dataclass +class _NullPlan: + """Every random number one :func:`null_model_und_sign` call uses. - Direct port of ``null_model_und_sign.m``'s ``wei_freq < 1`` (periodic - re-sort) branch — the default in modern MATLAB (``nargin('randperm')~=1`` - always true in any MATLAB version this codebase targets), so the - ``wei_freq==1`` exact-resort branch isn't ported. + ``seed`` is the rewiring's (None when the network is not rewired, or + numba is missing and the rewiring draws for itself), and ``choices`` + holds, per sign, the ``rng.choice`` draw of each weight-assignment pass. """ - if rng is None: - rng = np.random.default_rng() + rewire: bool + seed: int | None + choices: tuple[list[np.ndarray], list[np.ndarray]] + + +def _draw_null_plan(w: np.ndarray, bin_swaps: int, wei_freq: float, + rng: np.random.Generator) -> _NullPlan: + """Draw a null model's random numbers in the order the model uses them. + + Possible because none of the draws depends on the rewired network: the + rewiring takes one seed, and each weight-assignment pass draws a choice + whose size depends only on how many edges of that sign the *original* + network has (rewiring preserves the count). So the draws can be taken + first, in order, and the expensive rewiring run afterwards, on any + thread, with the result and the generator's final state unchanged. + """ + n = w.shape[0] + w = w.astype(float).copy() + np.fill_diagonal(w, 0.0) + ap = w > 0 + an = w < 0 + rewire = bool(np.count_nonzero(ap) < n * (n - 1)) + seed = None + if rewire and _HAVE_NUMBA and n >= 4: + seed = int(rng.integers(0, 2**31 - 1)) + elif rewire and n >= 4: + raise RuntimeError("_draw_null_plan needs numba for the rewiring") + + wei_period = round(1 / wei_freq) + choices: tuple[list[np.ndarray], list[np.ndarray]] = ([], []) + for draws, a_mask in zip(choices, (ap, an)): + m = int(np.count_nonzero(np.triu(a_mask))) + while m > 0: + batch = min(m, wei_period) + draws.append(rng.choice(m, size=batch, replace=False)) + m -= batch + return _NullPlan(rewire, seed, choices) + + +def _null_from_plan(w: np.ndarray, bin_swaps: int, wei_freq: float, + plan: _NullPlan | None, rng: np.random.Generator | None = None) -> np.ndarray: + """Build the null model, taking random numbers from ``plan`` or, with no + plan, straight from ``rng`` (the path without numba).""" n = w.shape[0] w = w.astype(float).copy() np.fill_diagonal(w, 0.0) @@ -207,7 +259,12 @@ def null_model_und_sign( an = w < 0 if np.count_nonzero(ap) < n * (n - 1): - w_r = randmio_und_signed(w, bin_swaps, rng=rng) + if plan is None: + w_r = randmio_und_signed(w, bin_swaps, rng=rng) + elif plan.seed is None: + w_r = w.copy() # n < 4: nothing is rewirable (see randmio_und_signed) + else: + w_r = _randmio_signed_seeded(w.copy(), bin_swaps, plan.seed) ap_r = w_r > 0 an_r = w_r < 0 else: @@ -217,7 +274,7 @@ def null_model_und_sign( w0 = np.zeros((n, n)) wei_period = round(1 / wei_freq) - for sign, a_mask, a_mask_r in ((1, ap, ap_r), (-1, an, an_r)): + for sign, a_mask, a_mask_r, planned in ((1, ap, ap_r, 0), (-1, an, an_r, 1)): if sign == 1: s = (w * a_mask).sum(axis=1) wv = np.sort(w[np.triu(a_mask)]) @@ -231,6 +288,7 @@ def null_model_und_sign( lij = n * j_idx + i_idx p = np.outer(s, s) + draws = iter(plan.choices[planned]) if plan is not None else None # ``i_idx``/``j_idx``/``lij``/``wv`` are kept as arrays rather than # lists: each pass drops only ``wei_period`` entries, so the loop runs @@ -242,7 +300,7 @@ def null_model_und_sign( while m > 0: batch = min(m, wei_period) p_lij = p.flat[lij] - r_idx = rng.choice(m, size=batch, replace=False) + r_idx = next(draws) if draws is not None else rng.choice(m, size=batch, replace=False) o = _stable_ranks(p_lij, r_idx) # These three stay Python-level: they are ``batch`` (<=10) items, @@ -282,67 +340,240 @@ def null_model_und_sign( return w0 + w0.T +def null_model_und_sign( + w: np.ndarray, + bin_swaps: int = 5, + wei_freq: float = 0.1, + rng: np.random.Generator | None = None, +) -> np.ndarray: + """Randomize an undirected network, preserving degree and (approximately) + strength distributions. + + Direct port of ``null_model_und_sign.m``'s ``wei_freq < 1`` (periodic + re-sort) branch — the default in modern MATLAB (``nargin('randperm')~=1`` + always true in any MATLAB version this codebase targets), so the + ``wei_freq==1`` exact-resort branch isn't ported. + """ + if rng is None: + rng = np.random.default_rng() + plan = _draw_null_plan(w, bin_swaps, wei_freq, rng) if _HAVE_NUMBA else None + return _null_from_plan(w, bin_swaps, wei_freq, plan, rng) + + +def map_null_models_und_sign( + w: np.ndarray, + count: int, + fn: Callable[[np.ndarray], T], + bin_swaps: int = 5, + wei_freq: float = 0.1, + rng: np.random.Generator | None = None, + workers: int = 1, +) -> list[T]: + """``[fn(null_model_und_sign(w, ...)) for _ in range(count)]``, on threads. + + Same null models, in the same order, and the same generator state + afterwards, as ``count`` sequential calls: every model's random numbers are + drawn first, in sequence (:func:`_draw_null_plan`), and only then are the + models built, up to ``workers`` at a time. The rewiring that dominates the + cost releases the GIL. ``fn`` reduces each model as soon as it is built, + so ``count`` full matrices are never held at once. + """ + if rng is None: + rng = np.random.default_rng() + if not _HAVE_NUMBA or workers <= 1 or count <= 1: + return [fn(null_model_und_sign(w, bin_swaps, wei_freq, rng)) for _ in range(count)] + plans = [_draw_null_plan(w, bin_swaps, wei_freq, rng) for _ in range(count)] + with ThreadPoolExecutor(max_workers=min(workers, count)) as ex: + return list(ex.map(lambda plan: fn(_null_from_plan(w, bin_swaps, wei_freq, plan)), plans)) + + # ── Edge-swap rewiring for small-worldness (randmio_und_v2.m / latmio_und_v2.m) ── # # Distinct from randmio_und_signed above: these pick two *existing edges* to # swap (by index into the nonzero lower-triangle) rather than four random # node indices, and have no notion of sign (MEA-NAP only calls them on -# already-nonnegative adjacency matrices). Both batch their random draws for -# speed, same trick as ``randmio_und_signed``'s ``next_quad``. - -# Consecutive rejected draws after which _valid_quad_stream concludes the graph -# has no valid quad at all. See the deviation note in its docstring. +# already-nonnegative adjacency matrices). +# +# The swap loop is sequential (each accepted swap changes the edge list the +# next draw reads), so it is compiled with numba like randmio_und_signed. Unlike +# that function, it does not switch to numba's own RNG: the random numbers are +# still drawn from the caller's Generator, in batches of _DRAW_BATCH, by the +# Python wrapper. When the compiled loop runs out of a batch it returns, the +# wrapper draws the next one, and the loop resumes where it stopped. That +# keeps the draws in exactly the order the earlier pure-Python version made +# them, so seeded runs give the same null models, and so the same SW/SWw/CC/PL, +# as before the loop was compiled. + +# Consecutive rejected draws after which the swap loop concludes the graph has +# no valid quad at all. See the deviation note in _rewire's docstring. _QUAD_SEARCH_BUDGET = 2_000_000 +# Random numbers drawn per refill: one batch of edge indices for e1, one for +# e2, or one batch of coin flips. +_DRAW_BATCH = 8192 + +# Why _edge_swap_core returned. +_SWAPS_DONE = 0 +_NEED_EDGE_DRAWS = 1 +_NEED_FLIP_DRAWS = 2 +_NO_VALID_QUAD = 3 + +# Slots of the core's ``state`` array, which carries its loop position across +# a return-and-refill: the outer iteration, the attempt within it, the read +# position in each batch, the consecutive-rejection count, and a quad found +# just before the flip batch ran out (pending flag + its two edge indices). +_ITER, _ATTEMPT, _POS_E, _POS_F, _REJECTED, _PENDING, _E1, _E2 = range(8) + -def _valid_quad_stream( - rng: np.random.Generator, k: int, i_arr: list[int], j_arr: list[int], batch_size: int = 8192, -) -> Iterator[tuple[int, int, int, int, int, int] | None]: - """Yields ``(e1, e2, a, b, c, d)``: two distinct edge indices into - ``i_arr``/``j_arr`` whose four endpoint nodes are all distinct. - - ``i_arr``/``j_arr`` are read fresh on each yield (not snapshotted), so - callers may mutate them between ``next()`` calls — required, since the - edge-flip step below does exactly that. - - **Deviation from MATLAB.** ``latmio_und_v2.m`` / ``randmio_und_v2.m`` wrap - this search in a bare ``while 1``, which never terminates on a graph where - no two edges have four distinct endpoints — a triangle, a star, any small - or degenerate component. MEA-NAP's ephys networks are large enough that - MATLAB never hit it, but a cell-type subnetwork can easily be three or four - cells, so the search is bounded here: after - ``_QUAD_SEARCH_BUDGET`` consecutive rejected draws this yields ``None``, - telling the caller to stop rewiring rather than spin forever. The budget is - far beyond what any graph that *does* admit a swap would need, so this only - changes behaviour in cases where MATLAB would hang. +def _edge_swap_core(r, i_arr, j_arr, d, lattice, iterations, max_attempts, budget, + e1_batch, e2_batch, flips, state): # pragma: no cover - jitted + """The rewiring loop shared by ``randmio_und_v2`` and ``latmio_und_v2``. + + Mutates ``r``, ``i_arr``, ``j_arr`` and ``state`` in place and returns one + of the ``_SWAPS_DONE``/``_NEED_*``/``_NO_VALID_QUAD`` codes. ``lattice`` + adds latmio's condition that a swap must not move edge weight onto more + distant pairs (``d``); with it off, ``d`` is never read. + """ + n_e = e1_batch.shape[0] + n_f = flips.shape[0] + it = state[_ITER] + attempt = state[_ATTEMPT] + pos_e = state[_POS_E] + pos_f = state[_POS_F] + rejected = state[_REJECTED] + status = _SWAPS_DONE + + while it < iterations: + swapped = False + while attempt <= max_attempts: + if state[_PENDING] == 0: + # Draw edge pairs until both edges' four endpoints are distinct. + found = False + while not found: + if pos_e >= n_e: + status = _NEED_EDGE_DRAWS + break + e1 = e1_batch[pos_e] + e2 = e2_batch[pos_e] + pos_e += 1 + if e1 != e2: + a = i_arr[e1] + b = j_arr[e1] + c = i_arr[e2] + dd = j_arr[e2] + if a != c and a != dd and b != c and b != dd: + rejected = 0 + found = True + continue + rejected += 1 + if rejected >= budget: + status = _NO_VALID_QUAD + break + if not found: + break + state[_PENDING] = 1 + state[_E1] = e1 + state[_E2] = e2 + if pos_f >= n_f: + status = _NEED_FLIP_DRAWS + break + flip = flips[pos_f] > 0.5 + pos_f += 1 + state[_PENDING] = 0 + + e1 = state[_E1] + e2 = state[_E2] + a = i_arr[e1] + b = j_arr[e1] + c = i_arr[e2] + dd = j_arr[e2] + if flip: + i_arr[e2] = dd + j_arr[e2] = c + c, dd = dd, c + # ``not (r[a, dd] or r[c, b])`` in the reference; a NaN counts as + # an edge there, and ``== 0`` is False for NaN, so they agree. + if r[a, dd] == 0 and r[c, b] == 0: + if not lattice or (d[a, b] * r[a, b] + d[c, dd] * r[c, dd] + >= d[a, dd] * r[a, b] + d[c, b] * r[c, dd]): + r_ab = r[a, b] + r_ba = r[b, a] + r_cd = r[c, dd] + r_dc = r[dd, c] + r[a, dd] = r_ab + r[a, b] = 0.0 + r[dd, a] = r_ba + r[b, a] = 0.0 + r[c, b] = r_cd + r[c, dd] = 0.0 + r[b, c] = r_dc + r[dd, c] = 0.0 + j_arr[e1] = dd + j_arr[e2] = b + swapped = True + break + attempt += 1 + if status != _SWAPS_DONE: + break + it += 1 + attempt = 0 + + state[_ITER] = it + state[_ATTEMPT] = attempt + state[_POS_E] = pos_e + state[_POS_F] = pos_f + state[_REJECTED] = rejected + return status + + +if _HAVE_NUMBA: + _edge_swap_core = njit(cache=True)(_edge_swap_core) + + +def _rewire(r: np.ndarray, d: np.ndarray, lattice: bool, iterations: int, + max_attempts: int, rng: np.random.Generator) -> None: + """Run the edge-swap loop on ``r`` in place, feeding it batches from ``rng``. + + **Deviation from MATLAB.** ``latmio_und_v2.m`` / ``randmio_und_v2.m`` search + for a swappable edge pair in a bare ``while 1``, which never terminates on + a graph where no two edges have four distinct endpoints — a triangle, a + star, any small or degenerate component. MEA-NAP's ephys networks are + large enough that MATLAB never hit it, but a cell-type subnetwork can + easily be three or four cells, so the search is bounded here: after + ``budget`` consecutive rejected draws, rewiring stops and the graph is + left as it is. The budget is far beyond what any graph that *does* admit a + swap would need, so this only changes behaviour where MATLAB would hang. """ + i_np, j_np = np.nonzero(np.tril(r)) + i_arr = np.ascontiguousarray(i_np, dtype=np.int64) + j_arr = np.ascontiguousarray(j_np, dtype=np.int64) + k = i_arr.shape[0] # A graph with any valid quad at all has at least ~2/k² of its draws # succeed, so scaling the budget with k² keeps the false-give-up # probability negligible while staying instant for the tiny graphs where # giving up is the correct answer. budget = min(_QUAD_SEARCH_BUDGET, max(10_000, 200 * k * k)) - rejected = 0 - while True: - e1_batch = rng.integers(0, k, size=batch_size) - e2_batch = rng.integers(0, k, size=batch_size) - for e1, e2 in zip(e1_batch.tolist(), e2_batch.tolist()): - if e1 != e2: - a, b = i_arr[e1], j_arr[e1] - c, d = i_arr[e2], j_arr[e2] - if a != c and a != d and b != c and b != d: - rejected = 0 - yield e1, e2, a, b, c, d - continue - rejected += 1 - if rejected >= budget: - yield None - rejected = 0 - - -def _coin_flip_stream(rng: np.random.Generator, batch_size: int = 8192) -> Iterator[bool]: + + # Both batches start "used up", so nothing is drawn before the loop asks: + # the same lazy order the generator-based version drew in. + e1_batch = np.zeros(_DRAW_BATCH, dtype=np.int64) + e2_batch = np.zeros(_DRAW_BATCH, dtype=np.int64) + flips = np.zeros(_DRAW_BATCH) + state = np.zeros(8, dtype=np.int64) + state[_POS_E] = _DRAW_BATCH + state[_POS_F] = _DRAW_BATCH while True: - for v in rng.random(batch_size).tolist(): - yield v > 0.5 + status = _edge_swap_core(r, i_arr, j_arr, d, lattice, iterations, max_attempts, + budget, e1_batch, e2_batch, flips, state) + if status == _NEED_EDGE_DRAWS: + e1_batch = rng.integers(0, k, size=_DRAW_BATCH) + e2_batch = rng.integers(0, k, size=_DRAW_BATCH) + state[_POS_E] = 0 + elif status == _NEED_FLIP_DRAWS: + flips = rng.random(_DRAW_BATCH) + state[_POS_F] = 0 + else: + return def randmio_und_v2( @@ -360,34 +591,12 @@ def randmio_und_v2( r = w.astype(float).copy() n = r.shape[0] - i_arr_np, j_arr_np = np.nonzero(np.tril(r)) - i_arr, j_arr = i_arr_np.tolist(), j_arr_np.tolist() - k = len(i_arr) + k = int(np.count_nonzero(np.tril(r))) if k < 2: return r max_attempts = round(n * k / (n * (n - 1))) - quads = _valid_quad_stream(rng, k, i_arr, j_arr) - flips = _coin_flip_stream(rng) - - for _ in range(iterations): - for _attempt in range(max_attempts + 1): - quad = next(quads) - if quad is None: - return r # no swappable edge pair exists — leave the graph as is - e1, e2, a, b, c, d = quad - if next(flips): - i_arr[e2], j_arr[e2] = d, c - c, d = d, c - if not (r[a, d] or r[c, b]): - r[a, d], r[a, b] = r[a, b], 0.0 - r[d, a], r[b, a] = r[b, a], 0.0 - r[c, b], r[c, d] = r[c, d], 0.0 - r[b, c], r[d, c] = r[d, c], 0.0 - j_arr[e1] = d - j_arr[e2] = b - break - + _rewire(r, np.zeros((1, 1)), False, iterations, max_attempts, rng) return r @@ -410,38 +619,10 @@ def latmio_und_v2( ind_rp = rng.permutation(n) r = w[np.ix_(ind_rp, ind_rp)].astype(float).copy() - i_arr_np, j_arr_np = np.nonzero(np.tril(r)) - i_arr, j_arr = i_arr_np.tolist(), j_arr_np.tolist() - k = len(i_arr) + k = int(np.count_nonzero(np.tril(r))) if k >= 2: max_attempts = round(n * k / (n * (n - 1) / 2)) - quads = _valid_quad_stream(rng, k, i_arr, j_arr) - flips = _coin_flip_stream(rng) - - for _ in range(iterations): - exhausted = False - for _attempt in range(max_attempts + 1): - quad = next(quads) - if quad is None: - exhausted = True # no swappable edge pair — stop rewiring - break - e1, e2, a, b, c, dd = quad - if next(flips): - i_arr[e2], j_arr[e2] = dd, c - c, dd = dd, c - if not (r[a, dd] or r[c, b]): - if (d[a, b] * r[a, b] + d[c, dd] * r[c, dd]) >= ( - d[a, dd] * r[a, b] + d[c, b] * r[c, dd] - ): - r[a, dd], r[a, b] = r[a, b], 0.0 - r[dd, a], r[b, a] = r[b, a], 0.0 - r[c, b], r[c, dd] = r[c, dd], 0.0 - r[b, c], r[dd, c] = r[dd, c], 0.0 - j_arr[e1] = dd - j_arr[e2] = b - break - if exhausted: - break + _rewire(r, np.ascontiguousarray(d, dtype=float), True, iterations, max_attempts, rng) ind_rp_reverse = np.argsort(ind_rp) return r[np.ix_(ind_rp_reverse, ind_rp_reverse)] diff --git a/src/meanap/pipeline/parallel.py b/src/meanap/pipeline/parallel.py index ffc158ce..0fe89935 100644 --- a/src/meanap/pipeline/parallel.py +++ b/src/meanap/pipeline/parallel.py @@ -155,6 +155,30 @@ def pin_blas_threads(threads: str = "1") -> None: os.environ.update(worker_env(threads)) +def compute_threads() -> int: + """Threads this process may use for its own parallel loops (numba + ``prange``, NMF fits). + + Inside a :func:`map_recordings` worker that is the BLAS share the worker + was pinned to, so N workers x their threads still adds up to the machine. + Outside a pool it is every physical core, unless ``OMP_NUM_THREADS`` says + otherwise. + """ + pinned = os.environ.get("OMP_NUM_THREADS", "") + if pinned.isdigit() and int(pinned) > 0: + return int(pinned) + return physical_cores() + + +def numba_threads() -> int: + """:func:`compute_threads`, capped at the size of numba's thread pool.""" + try: + import numba + except Exception: # pragma: no cover - numba optional + return 1 + return max(1, min(compute_threads(), numba.config.NUMBA_NUM_THREADS)) + + #: Cached across the session — whether ``__main__`` is guarded cannot change. _SPAWN_USABLE: Optional[bool] = None diff --git a/src/meanap/pipeline/step4.py b/src/meanap/pipeline/step4.py index d8872ede..4e0442a4 100644 --- a/src/meanap/pipeline/step4.py +++ b/src/meanap/pipeline/step4.py @@ -673,6 +673,7 @@ def _step4_compute_one( params.nmf_downsample_freq, fs, include_nmf_components=params.include_nmf_components, rng=dimensionality_rng, + plateau_gain=params.nmf_plateau_gain, ) except Exception as e: logs.append(f" [{rec.filename}] WARNING: could not compute NMF components: {e}") diff --git a/src/meanap/pipeline/sttc.py b/src/meanap/pipeline/sttc.py index 1fdf7119..9cec7157 100644 --- a/src/meanap/pipeline/sttc.py +++ b/src/meanap/pipeline/sttc.py @@ -7,8 +7,12 @@ from __future__ import annotations +import warnings + import numpy as np +from meanap.pipeline.parallel import numba_threads + # ``run_P`` (below) is a monotonic two-pointer whose result is *order-dependent* # on the input spike trains — it does not sort. For spike-detection data the # trains are already sorted, but CAT-NAP peak times can arrive unsorted (see @@ -17,7 +21,7 @@ # with numba when available (same optional pattern as ``null_models.py``); the # pure-Python fallback is identical, just slower. try: - from numba import njit + from numba import njit, prange, set_num_threads _HAVE_NUMBA = True except Exception: # pragma: no cover - numba optional / version-gated @@ -47,6 +51,69 @@ def _run_p_impl(t1: np.ndarray, t2: np.ndarray, dt: float) -> int: _run_p_jit = njit(cache=True)(_run_p_impl) if _HAVE_NUMBA else _run_p_impl +def _sttc_row_impl(flat: np.ndarray, starts: np.ndarray, tiled: np.ndarray, + dt: float, out: np.ndarray, i: int) -> None: + """Fill ``out[i, i+1:]`` with the STTC of train ``i`` against each later one. + + Train ``i`` is ``flat[starts[i]:starts[i + 1]]`` and ``tiled[i]`` its + normalised ``run_T``. The arithmetic is :func:`_sttc_from_tiled`'s, + written out so a whole sweep is one compiled call. Pairs with an empty + train are left as they are (NaN, which ``get_sttc`` zeroes). + """ + n = starts.shape[0] - 1 + n1 = starts[i + 1] - starts[i] + if n1 == 0: + return + t1 = flat[starts[i]:starts[i + 1]] + ta = tiled[i] + for j in range(i + 1, n): + n2 = starts[j + 1] - starts[j] + if n2 == 0: + continue + t2 = flat[starts[j]:starts[j + 1]] + tb = tiled[j] + pa = _run_p_jit(t1, t2, dt) / n1 + pb = _run_p_jit(t2, t1, dt) / n2 + out[i, j] = 0.5 * (pa - tb) / (1 - tb * pa) + 0.5 * (pb - ta) / (1 - ta * pb) + + +def _sttc_upper_impl(flat: np.ndarray, starts: np.ndarray, tiled: np.ndarray, + dt: float, out: np.ndarray) -> None: + """Fill the upper triangle of ``out``, rows in parallel. + + Row ``i`` has ``n - 1 - i`` pairs, so rows are taken in pairs from both + ends (``f`` and ``n - 1 - f``), giving every parallel step the same + ``n - 1`` pairs. Each pair is computed exactly as in the serial loop and + written to its own cell, so the result does not depend on the thread + count. + """ + n = starts.shape[0] - 1 + for f in _prange((n + 1) // 2): + _sttc_row(flat, starts, tiled, dt, out, f) + g = n - 1 - f + if g != f: + _sttc_row(flat, starts, tiled, dt, out, g) + + +# ``error_model="numpy"``: a zero denominator gives inf/NaN, as it does for +# the NumPy scalars in :func:`_sttc_from_tiled`, instead of raising. +if _HAVE_NUMBA: + _prange = prange + _sttc_row = njit(cache=True, error_model="numpy")(_sttc_row_impl) + _sttc_upper = njit(cache=True, error_model="numpy", parallel=True)(_sttc_upper_impl) + # The same loop compiled without ``parallel``, where ``prange`` is a plain + # range: what runs if no threading layer can be loaded (see ``get_sttc``). + _sttc_upper_serial = njit(cache=True, error_model="numpy")(_sttc_upper_impl) +else: # pragma: no cover - numba optional + _prange = range + _sttc_row = _sttc_row_impl + _sttc_upper = _sttc_upper_serial = _sttc_upper_impl + +#: Set once the parallel kernel has failed to start, so later calls go straight +#: to the serial one. +_PARALLEL_BROKEN = False + + def _run_t(dt: float, t_start: float, t_end: float, times: np.ndarray) -> float: """Fraction-of-recording (unnormalised) tiled by ``times`` +/- ``dt``. @@ -152,14 +219,36 @@ def get_sttc( tiled = [(_run_t(dt, 0.0, duration_s, tr) / t_span) if tr.size else 0.0 for tr in times] - for i in range(n_channels): - ti = times[i] - if ti.size == 0: - continue - for j in range(i + 1, n_channels): - coef = _sttc_from_tiled(ti, times[j], dt, tiled[i], tiled[j]) - adj_m[i, j] = coef - adj_m[j, i] = coef + # The pair sweep itself costs O(n_channels x total spikes): every pair walks + # both trains. That is negligible on a 64-electrode MEA and ~2s per matrix + # for 300 sorted units (so ~20 min of step 3 at 201 matrices x 3 lags), + # which is why it runs as one compiled call, rows across cores. Trains go + # in as one flat float64 array; spike times are float64 everywhere in the + # pipeline, so this is the arithmetic the per-pair path did. + starts = np.zeros(n_channels + 1, dtype=np.int64) + np.cumsum([tr.size for tr in times], out=starts[1:]) + flat = (np.concatenate(times).astype(np.float64, copy=False) if n_channels + else np.empty(0)) + args = (flat, starts, np.asarray(tiled, dtype=np.float64), float(dt), adj_m) + global _PARALLEL_BROKEN + if _HAVE_NUMBA and not _PARALLEL_BROKEN: + try: + set_num_threads(numba_threads()) + _sttc_upper(*args) + except Exception as e: + # numba loads its threading layer on the first parallel call, and + # a build that is missing one (a frozen app, an unusual install) + # fails here rather than at import. The serial kernel gives the + # same matrix, just on one core. + if "threading layer" not in str(e).lower(): + raise + _PARALLEL_BROKEN = True + warnings.warn(f"STTC runs on one core: numba could not start threads ({e})") + _sttc_upper_serial(*args) + else: + _sttc_upper_serial(*args) + lower = np.tril_indices(n_channels, -1) + adj_m[lower] = adj_m.T[lower] adj_m[adj_m < 0] = 0.0 adj_m[np.isnan(adj_m)] = 0.0 diff --git a/version.txt b/version.txt index 63e799cf..141f2e80 100644 --- a/version.txt +++ b/version.txt @@ -1 +1 @@ -1.14.1 +1.15.0 diff --git a/versions.json b/versions.json index 668c07f2..5d7e6edc 100644 --- a/versions.json +++ b/versions.json @@ -1,5 +1,5 @@ { "_comment": "Per-subsystem versions. MEA-NAP's own version lives in version.txt, which MATLAB reads and the update check compares against GitHub; it is not duplicated here.", - "catnap": "1.10.2", + "catnap": "1.11.0", "meastim": "0.5.0" }