From 97b242ddcc1f6d62c9943174c21b183277ff0131 Mon Sep 17 00:00:00 2001 From: "P. L. Lim" <2090236+pllim@users.noreply.github.com> Date: Thu, 2 Jul 2026 18:41:57 -0400 Subject: [PATCH 1/2] Shared memory: sharing outputs no work but only inputs worked earlier --- acstools/utils_findsat_mrt.py | 89 +++++++++++++++++++++++++++++++---- 1 file changed, 79 insertions(+), 10 deletions(-) diff --git a/acstools/utils_findsat_mrt.py b/acstools/utils_findsat_mrt.py index 70bda8e..67518e1 100644 --- a/acstools/utils_findsat_mrt.py +++ b/acstools/utils_findsat_mrt.py @@ -6,6 +6,8 @@ import time import warnings from itertools import repeat +from multiprocessing import shared_memory +from multiprocessing.managers import SharedMemoryManager import numpy as np from astropy.convolution import convolve, Gaussian2DKernel @@ -1012,6 +1014,50 @@ def rot_med(image, angle, return_length): return _rot(image, angle, return_length, np.nanmedian) +def _rot_sum_mp(shm_padded_name, shm_radon_name, shm_lengths_name, + image_shape, image_dtype, n_angles, i_angle, angle): + """rot_sum for multiprocessing""" + # Grab input array from shared memory + shm_padded = shared_memory.SharedMemory(name=shm_padded_name) + image = np.ndarray( + image_shape, dtype=image_dtype, buffer=shm_padded.buf) + + med_arr, length = _rot(image, angle, True, np.nansum) + + # Put result in shared memory + result_shape = (image_shape[0], n_angles) + shm_radon = shared_memory.SharedMemory(name=shm_radon_name) + radon_image = np.ndarray( + result_shape, dtype=image_dtype, buffer=shm_radon.buf) + shm_lengths = shared_memory.SharedMemory(name=shm_lengths_name) + lengths = np.ndarray( + result_shape, dtype=image_dtype, buffer=shm_lengths.buf) + radon_image[:, i_angle] = med_arr + lengths[:, i_angle] = length + + +def _rot_med_mp(shm_padded_name, shm_radon_name, shm_lengths_name, + image_shape, image_dtype, n_angles, i_angle, angle): + """rot_med for multiprocessing""" + # Grab input array from shared memory + shm_padded = shared_memory.SharedMemory(name=shm_padded_name) + image = np.ndarray( + image_shape, dtype=image_dtype, buffer=shm_padded.buf) + + med_arr, length = _rot(image, angle, True, np.nanmedian) + + # Put result in shared memory + result_shape = (image_shape[0], n_angles) + shm_radon = shared_memory.SharedMemory(name=shm_radon_name) + radon_image = np.ndarray( + result_shape, dtype=image_dtype, buffer=shm_radon.buf) + shm_lengths = shared_memory.SharedMemory(name=shm_lengths_name) + lengths = np.ndarray( + result_shape, dtype=image_dtype, buffer=shm_lengths.buf) + radon_image[:, i_angle] = med_arr + lengths[:, i_angle] = length + + # TODO: If radon performance is improved upstream, we should just use # the version in scikit-image and remove this one. See # https://github.com/scikit-image/scikit-image/issues/3118 @@ -1090,9 +1136,11 @@ def radon(image, theta=None, circle=False, *, preserve_range=False, if median is True: statname = "median" statfunc = rot_med + statfunc_mp = _rot_med_mp else: statname = "standard" statfunc = rot_sum + statfunc_mp = _rot_sum_mp LOG.info('Calculating %s Radon Transform with %d processes', statname, processes) image = convert_to_float(image, preserve_range) @@ -1132,7 +1180,8 @@ def radon(image, theta=None, circle=False, *, preserve_range=False, theta = np.arange(180) angles = np.deg2rad(theta) - radon_image = np.empty((padded_image.shape[0], len(theta)), dtype=image.dtype) + n_angles = len(angles) + radon_image = np.empty((padded_image.shape[0], n_angles), dtype=image.dtype) radon_image[:] = np.nan lengths = np.copy(radon_image) @@ -1149,19 +1198,39 @@ def radon(image, theta=None, circle=False, *, preserve_range=False, else: # Splitting calculation up among many processes to speed up. Each # thread rotates and sums/medians at a specific angle. - collected_results = {} - - with concurrent.futures.ProcessPoolExecutor(max_workers=processes) as p, warnings.catch_warnings(): + with concurrent.futures.ProcessPoolExecutor(max_workers=processes) as p, \ + SharedMemoryManager() as smm, \ + warnings.catch_warnings(): # suppressing this warning as it's inconsequential and expected warnings.filterwarnings('ignore', message='All-NaN slice encountered') - for angle, result in zip(angles, p.map(statfunc, repeat(padded_image), angles, repeat(True))): - collected_results[angle] = result + # Shared memory for padded_image + shm_padded = smm.SharedMemory(padded_image.nbytes) + mparr_padded = np.ndarray( + padded_image.shape, dtype=padded_image.dtype, buffer=shm_padded.buf) + mparr_padded[:] = padded_image[:] - for i, angle in enumerate(angles): - result = collected_results[angle] - radon_image[:, i] = result[0] - lengths[:, i] = result[1] + # Shared memory for radon_image + shm_radon = smm.SharedMemory(radon_image.nbytes) + mparr_radon = np.ndarray( + radon_image.shape, dtype=radon_image.dtype, buffer=shm_radon.buf) + mparr_radon[:] = radon_image[:] + + # Shared memory for lengths + shm_lengths = smm.SharedMemory(lengths.nbytes) + mparr_lengths = np.ndarray( + lengths.shape, dtype=lengths.dtype, buffer=shm_lengths.buf) + mparr_lengths = lengths[:] + + for i, angle in enumerate(angles): + future = p.submit( + statfunc_mp, shm_padded.name, shm_radon.name, shm_lengths.name, + padded_image.shape, padded_image.dtype, n_angles, i, angle) + _ = future.result() + + # Copy from shared memory to outputs + radon_image[:] = mparr_radon[:] + lengths[:] = mparr_lengths[:] if return_length is True: return radon_image, lengths From e386098ac8e96087f0042815be52eb1b096606de Mon Sep 17 00:00:00 2001 From: "P. L. Lim" <2090236+pllim@users.noreply.github.com> Date: Thu, 2 Jul 2026 19:01:57 -0400 Subject: [PATCH 2/2] Revert shared mem for outputs --- acstools/utils_findsat_mrt.py | 73 +++++++++-------------------------- 1 file changed, 19 insertions(+), 54 deletions(-) diff --git a/acstools/utils_findsat_mrt.py b/acstools/utils_findsat_mrt.py index 67518e1..c172bc9 100644 --- a/acstools/utils_findsat_mrt.py +++ b/acstools/utils_findsat_mrt.py @@ -1014,48 +1014,22 @@ def rot_med(image, angle, return_length): return _rot(image, angle, return_length, np.nanmedian) -def _rot_sum_mp(shm_padded_name, shm_radon_name, shm_lengths_name, - image_shape, image_dtype, n_angles, i_angle, angle): +def _rot_sum_mp(shm_padded_name, image_shape, image_dtype, angle): """rot_sum for multiprocessing""" # Grab input array from shared memory - shm_padded = shared_memory.SharedMemory(name=shm_padded_name) + shm_padded = shared_memory.SharedMemory(name=shm_padded_name, create=False) image = np.ndarray( image_shape, dtype=image_dtype, buffer=shm_padded.buf) + return _rot(image, angle, True, np.nansum) - med_arr, length = _rot(image, angle, True, np.nansum) - # Put result in shared memory - result_shape = (image_shape[0], n_angles) - shm_radon = shared_memory.SharedMemory(name=shm_radon_name) - radon_image = np.ndarray( - result_shape, dtype=image_dtype, buffer=shm_radon.buf) - shm_lengths = shared_memory.SharedMemory(name=shm_lengths_name) - lengths = np.ndarray( - result_shape, dtype=image_dtype, buffer=shm_lengths.buf) - radon_image[:, i_angle] = med_arr - lengths[:, i_angle] = length - - -def _rot_med_mp(shm_padded_name, shm_radon_name, shm_lengths_name, - image_shape, image_dtype, n_angles, i_angle, angle): +def _rot_med_mp(shm_padded_name, image_shape, image_dtype, angle): """rot_med for multiprocessing""" # Grab input array from shared memory - shm_padded = shared_memory.SharedMemory(name=shm_padded_name) + shm_padded = shared_memory.SharedMemory(name=shm_padded_name, create=False) image = np.ndarray( image_shape, dtype=image_dtype, buffer=shm_padded.buf) - - med_arr, length = _rot(image, angle, True, np.nanmedian) - - # Put result in shared memory - result_shape = (image_shape[0], n_angles) - shm_radon = shared_memory.SharedMemory(name=shm_radon_name) - radon_image = np.ndarray( - result_shape, dtype=image_dtype, buffer=shm_radon.buf) - shm_lengths = shared_memory.SharedMemory(name=shm_lengths_name) - lengths = np.ndarray( - result_shape, dtype=image_dtype, buffer=shm_lengths.buf) - radon_image[:, i_angle] = med_arr - lengths[:, i_angle] = length + return _rot(image, angle, True, np.nanmedian) # TODO: If radon performance is improved upstream, we should just use @@ -1197,7 +1171,7 @@ def radon(image, theta=None, circle=False, *, preserve_range=False, else: # Splitting calculation up among many processes to speed up. Each - # thread rotates and sums/medians at a specific angle. + # process rotates and sums/medians at a specific angle. with concurrent.futures.ProcessPoolExecutor(max_workers=processes) as p, \ SharedMemoryManager() as smm, \ warnings.catch_warnings(): @@ -1210,27 +1184,18 @@ def radon(image, theta=None, circle=False, *, preserve_range=False, padded_image.shape, dtype=padded_image.dtype, buffer=shm_padded.buf) mparr_padded[:] = padded_image[:] - # Shared memory for radon_image - shm_radon = smm.SharedMemory(radon_image.nbytes) - mparr_radon = np.ndarray( - radon_image.shape, dtype=radon_image.dtype, buffer=shm_radon.buf) - mparr_radon[:] = radon_image[:] - - # Shared memory for lengths - shm_lengths = smm.SharedMemory(lengths.nbytes) - mparr_lengths = np.ndarray( - lengths.shape, dtype=lengths.dtype, buffer=shm_lengths.buf) - mparr_lengths = lengths[:] - - for i, angle in enumerate(angles): - future = p.submit( - statfunc_mp, shm_padded.name, shm_radon.name, shm_lengths.name, - padded_image.shape, padded_image.dtype, n_angles, i, angle) - _ = future.result() - - # Copy from shared memory to outputs - radon_image[:] = mparr_radon[:] - lengths[:] = mparr_lengths[:] + for i, angle, result in zip( + range(n_angles), + angles, p.map( + statfunc_mp, + repeat(shm_padded.name), + repeat(padded_image.shape), + repeat(padded_image.dtype), + angles, + ) + ): + radon_image[:, i] = result[0] + lengths[:, i] = result[1] if return_length is True: return radon_image, lengths