diff --git a/acstools/utils_findsat_mrt.py b/acstools/utils_findsat_mrt.py index 70bda8e..c172bc9 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,24 @@ def rot_med(image, angle, return_length): return _rot(image, angle, return_length, np.nanmedian) +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, create=False) + image = np.ndarray( + image_shape, dtype=image_dtype, buffer=shm_padded.buf) + return _rot(image, angle, True, np.nansum) + + +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, create=False) + image = np.ndarray( + image_shape, dtype=image_dtype, buffer=shm_padded.buf) + return _rot(image, angle, True, np.nanmedian) + + # 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 +1110,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 +1154,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) @@ -1148,20 +1171,31 @@ 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(): + # process rotates and sums/medians at a specific angle. + 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 - - for i, angle in enumerate(angles): - result = collected_results[angle] - radon_image[:, i] = result[0] - lengths[:, i] = result[1] + # 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, 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