From 07a24a4cf2e4a82c70330b94c0ea6641ae27397d Mon Sep 17 00:00:00 2001 From: Geert Jan Talens Date: Thu, 17 Jul 2025 14:33:14 +0100 Subject: [PATCH 1/7] Fixes warning regarding the use of os.fork when using jax. --- nuance/periodic_search.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/nuance/periodic_search.py b/nuance/periodic_search.py index 39540ba..25e865b 100644 --- a/nuance/periodic_search.py +++ b/nuance/periodic_search.py @@ -33,7 +33,7 @@ def periodic_search(epochs, durations, ls, snr_f, progress=True): callable Function that computes the SNR and parameters for each period. """ - global fold_f + fold_f = _fold_ll(epochs, *ls) def _progress(x, **kwargs): @@ -43,9 +43,10 @@ def function(periods): snr = np.zeros(len(periods)) params = np.zeros((len(periods), 3)) - with mp.Pool() as pool: + ctx = mp.get_context('spawn') # Can't use fork with jax. + with ctx.Pool() as pool: for p, (epoch, duration_i, period) in enumerate( - _progress(pool.imap(_solve, periods), total=len(periods)) + _progress(pool.starmap(_solve, [(period, fold_f) for period in periods]), total=len(periods)) ): Dj = durations[duration_i] snr[p], params[p] = float(snr_f(epoch, Dj, period)), (epoch, Dj, period) @@ -84,7 +85,7 @@ def fun(period): return fun -def _solve(period): +def _solve(period, fold_f): phase, lls = fold_f(period) epoch_i, duration_i = np.unravel_index(np.argmax(lls), lls.shape) epoch = phase[epoch_i] * period From 11dc71b25b60b0606488db775c5a803f2c065950 Mon Sep 17 00:00:00 2001 From: Geert Jan Talens Date: Thu, 17 Jul 2025 15:15:17 +0100 Subject: [PATCH 2/7] Exposes some multiprocessing parameters. --- nuance/periodic_search.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/nuance/periodic_search.py b/nuance/periodic_search.py index 25e865b..17cb826 100644 --- a/nuance/periodic_search.py +++ b/nuance/periodic_search.py @@ -39,14 +39,14 @@ def periodic_search(epochs, durations, ls, snr_f, progress=True): def _progress(x, **kwargs): return tqdm(x, **kwargs) if progress else x - def function(periods): + def function(periods, processes=None, chunksize=500): snr = np.zeros(len(periods)) params = np.zeros((len(periods), 3)) ctx = mp.get_context('spawn') # Can't use fork with jax. - with ctx.Pool() as pool: + with ctx.Pool(processes=processes) as pool: for p, (epoch, duration_i, period) in enumerate( - _progress(pool.starmap(_solve, [(period, fold_f) for period in periods]), total=len(periods)) + _progress(pool.starmap(_solve, [(period, fold_f) for period in periods], chunksize=chunksize), total=len(periods)) ): Dj = durations[duration_i] snr[p], params[p] = float(snr_f(epoch, Dj, period)), (epoch, Dj, period) From 148182b3c2b8a6d38f6e453fa4283d0950c0c030 Mon Sep 17 00:00:00 2001 From: Geert Jan Talens Date: Fri, 18 Jul 2025 15:39:34 +0100 Subject: [PATCH 3/7] multiprocessing good practice change. --- nuance/periodic_search.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/nuance/periodic_search.py b/nuance/periodic_search.py index 17cb826..0b19af6 100644 --- a/nuance/periodic_search.py +++ b/nuance/periodic_search.py @@ -90,3 +90,11 @@ def _solve(period, fold_f): epoch_i, duration_i = np.unravel_index(np.argmax(lls), lls.shape) epoch = phase[epoch_i] * period return epoch, duration_i, period + + +def main(): + return + + +if __name__ == '__main__': + main() From fffae8a3294bd316699faf17b13996f25952f0a4 Mon Sep 17 00:00:00 2001 From: Geert Jan Talens Date: Fri, 25 Jul 2025 21:09:13 +0100 Subject: [PATCH 4/7] Minor changes to _fold_ll --- nuance/periodic_search.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/nuance/periodic_search.py b/nuance/periodic_search.py index 0b19af6..2ac1fc3 100644 --- a/nuance/periodic_search.py +++ b/nuance/periodic_search.py @@ -62,11 +62,10 @@ def _fold_ll(epochs, lls, z, vz): f_dz2 = core.nearest_neighbors(epochs, vz) def _fold(times): - lls = np.array([f_ll(time) for time in times]) - zs = np.array([f_z(time) for time in times]) - vzs = np.array([f_dz2(time) for time in times]) + lls = f_ll(times) + zs = f_z(times) + vzs = f_dz2(times) - P1 = np.sum(lls, 0) vZ = 1 / np.sum(1 / vzs, 0) Z = vZ * np.sum(zs / vzs, 0) P1 = np.sum(lls, 0) From dd608d0f2805daa2d57f9efa70883670cc776951 Mon Sep 17 00:00:00 2001 From: Geert Jan Talens Date: Fri, 25 Jul 2025 22:56:16 +0100 Subject: [PATCH 5/7] Use jax.vmap for computing SNR values. --- nuance/periodic_search.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/nuance/periodic_search.py b/nuance/periodic_search.py index 2ac1fc3..b87a18c 100644 --- a/nuance/periodic_search.py +++ b/nuance/periodic_search.py @@ -5,10 +5,11 @@ """ import multiprocess as mp +import jax import numpy as np from tqdm.auto import tqdm -from nuance import core +from nuance import DEVICES_COUNT, core from nuance.utils import interp_split_times @@ -39,7 +40,7 @@ def periodic_search(epochs, durations, ls, snr_f, progress=True): def _progress(x, **kwargs): return tqdm(x, **kwargs) if progress else x - def function(periods, processes=None, chunksize=500): + def function(periods, processes=DEVICES_COUNT, chunksize=500, batch_size=DEVICES_COUNT): snr = np.zeros(len(periods)) params = np.zeros((len(periods), 3)) @@ -49,7 +50,14 @@ def function(periods, processes=None, chunksize=500): _progress(pool.starmap(_solve, [(period, fold_f) for period in periods], chunksize=chunksize), total=len(periods)) ): Dj = durations[duration_i] - snr[p], params[p] = float(snr_f(epoch, Dj, period)), (epoch, Dj, period) + params[p] = (epoch, Dj, period) + + # Use jax.vmap to get the SNR at each period. + snr_vmap = jax.vmap(snr_f, in_axes=(0, 0, 0)) + for i in _progress(range(0, len(periods), batch_size), unit_scale=batch_size): + imin = i + imax = i + batch_size + snr[imin:imax] = snr_vmap(params[imin:imax, 0], params[imin:imax, 1], params[imin:imax, 2]) return snr, params From 09156ca8eeeb24001a68ff02be1b0a59d0193836 Mon Sep 17 00:00:00 2001 From: Geert Jan Talens Date: Fri, 25 Jul 2025 23:13:07 +0100 Subject: [PATCH 6/7] Rewrite multiprocessing so no chunksize is needed. --- nuance/periodic_search.py | 36 +++++++++++++++++++++++++----------- 1 file changed, 25 insertions(+), 11 deletions(-) diff --git a/nuance/periodic_search.py b/nuance/periodic_search.py index b87a18c..88bdeb5 100644 --- a/nuance/periodic_search.py +++ b/nuance/periodic_search.py @@ -5,6 +5,7 @@ """ import multiprocess as mp +from functools import partial import jax import numpy as np from tqdm.auto import tqdm @@ -40,17 +41,21 @@ def periodic_search(epochs, durations, ls, snr_f, progress=True): def _progress(x, **kwargs): return tqdm(x, **kwargs) if progress else x - def function(periods, processes=DEVICES_COUNT, chunksize=500, batch_size=DEVICES_COUNT): + def function(periods, processes=DEVICES_COUNT, batch_size=DEVICES_COUNT): snr = np.zeros(len(periods)) params = np.zeros((len(periods), 3)) + # Use multiprocessing to get the optimal epoch and duration at each period. + solve_f = partial(_solve, fold_f) ctx = mp.get_context('spawn') # Can't use fork with jax. with ctx.Pool(processes=processes) as pool: - for p, (epoch, duration_i, period) in enumerate( - _progress(pool.starmap(_solve, [(period, fold_f) for period in periods], chunksize=chunksize), total=len(periods)) - ): - Dj = durations[duration_i] - params[p] = (epoch, Dj, period) + period_chunks = [periods[i::processes] for i in range(processes)] + + for i, result in enumerate(_progress(pool.imap(solve_f, period_chunks), total=processes)): + epochs_chunk, duration_idx_chunk, periods_chunk = result + params[i::processes, 0] = epochs_chunk + params[i::processes, 1] = durations[duration_idx_chunk] + params[i::processes, 2] = periods_chunk # Use jax.vmap to get the SNR at each period. snr_vmap = jax.vmap(snr_f, in_axes=(0, 0, 0)) @@ -92,11 +97,20 @@ def fun(period): return fun -def _solve(period, fold_f): - phase, lls = fold_f(period) - epoch_i, duration_i = np.unravel_index(np.argmax(lls), lls.shape) - epoch = phase[epoch_i] * period - return epoch, duration_i, period +def _solve(fold_f, periods): + + epochs = np.zeros_like(periods) + duration_idx = np.zeros_like(periods, dtype='int') + + for i, period in enumerate(periods): + phase, lls = fold_f(period) + epoch_i, duration_i = np.unravel_index(np.argmax(lls), lls.shape) + epoch = phase[epoch_i] * period + + epochs[i] = epoch + duration_idx[i] = duration_i + + return epochs, duration_idx, periods def main(): From 80a5cbab39cbbbaf9efea8d85d2e3bb49319964c Mon Sep 17 00:00:00 2001 From: Geert Jan Talens Date: Mon, 4 Aug 2025 11:25:09 +0100 Subject: [PATCH 7/7] pmap is faster for snr computation. --- nuance/periodic_search.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nuance/periodic_search.py b/nuance/periodic_search.py index 88bdeb5..08e88f6 100644 --- a/nuance/periodic_search.py +++ b/nuance/periodic_search.py @@ -58,7 +58,7 @@ def function(periods, processes=DEVICES_COUNT, batch_size=DEVICES_COUNT): params[i::processes, 2] = periods_chunk # Use jax.vmap to get the SNR at each period. - snr_vmap = jax.vmap(snr_f, in_axes=(0, 0, 0)) + snr_vmap = jax.pmap(snr_f, in_axes=(0, 0, 0)) for i in _progress(range(0, len(periods), batch_size), unit_scale=batch_size): imin = i imax = i + batch_size