From ebdd83d871cb78d94f1ac3c1c03fbd8840f0619f Mon Sep 17 00:00:00 2001 From: Sam Nystrom Date: Fri, 16 Oct 2026 18:47:06 -0400 Subject: [PATCH] Add nproc argument to Coherence.signif_test If the caller is in a process pool, they may wish to disable or limit parallelization in pyleoclim. --- pyleoclim/core/coherences.py | 23 ++++++++++++++++++----- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/pyleoclim/core/coherences.py b/pyleoclim/core/coherences.py index d4cc8e06..a807cede 100644 --- a/pyleoclim/core/coherences.py +++ b/pyleoclim/core/coherences.py @@ -51,9 +51,9 @@ def _run_global_coherence(args): @contextmanager -def _get_process_pool(): +def _get_process_pool(nproc=None): ctx = multiprocessing.get_context("spawn") - with ProcessPoolExecutor(mp_context=ctx) as executor: + with ProcessPoolExecutor(mp_context=ctx, max_workers=nproc) as executor: yield executor @@ -753,6 +753,7 @@ def signif_test( qs=[0.95], settings=None, mute_pbar=False, + nproc=None, ): """Significance testing for Coherence objects @@ -788,6 +789,12 @@ def signif_test( Mute the progress bar. The default is False. + nproc : int | None, optional + + Number of worker processes for the surrogate coherences. None (default) uses + one per available CPU. 1 runs serially in the calling process with no pool + startup cost. + Returns ------- new : pyleoclim.core.coherence.Coherence @@ -888,16 +895,22 @@ def signif_test( ] # Perform wavelet coherence calculations in parallel - with _get_process_pool() as executor: - results = list( + def get_results(map_fn): + return list( tqdm( - executor.map(_run_wavelet_coherence, args), + map_fn(_run_wavelet_coherence, args), total=number, desc="Performing wavelet coherence on surrogate pairs", disable=mute_pbar, ) ) + if nproc == 1: + results = get_results(map) + else: + with _get_process_pool(nproc) as executor: + results = get_results(executor.map) + # Split results into wtcs and xwts wtcs = [result.wtc for result in results] xwts = [result.xwt for result in results]