cube_correct_nan updated with multiprocessing
#! /usr/bin/env python
-
+# check123
"""
Module with cosmetics procedures. Contains the function for bad pixel fixing.
Also functions for cropping cubes.
@@ -21,6 +21,9 @@ from astropy.stats import sigma_clipped_stats
from ..config.utils_conf import pool_map, iterable
from ..stats import sigma_filter
from ..var import frame_center, get_square
+from multiprocessing import Process
+import multiprocessing
+from multiprocessing import shared_memory, set_start_method
def cube_crop_frames(array, size, xy=None, force=False, verbose=True,
@@ -346,11 +349,41 @@ def cube_correct_nan(cube, neighbor_box=3, min_neighbors=3, verbose=False,
print(msg.format(zz, nnanpix))
else:
if verbose:
- msg = "Correcting NaNs in multiprocessing..."
+ msg = "Correcting NaNs in multiprocessing using ADACS' approach..."
print(msg)
- res = pool_map(nproc, nan_corr_2d, iterable(obj_tmp), neighbor_box,
- min_neighbors, half_res_y, verbose, False)
- obj_tmp = np.array(res, dtype=object)
+ #creation of shared memory blob
+ shm = shared_memory.SharedMemory(create=True, size=obj_tmp.nbytes)
+ #creating an array object similar to obj_tmp and occupying shared memory.
+ obj_tmp_shared=np.ndarray(obj_tmp.shape, dtype=obj_tmp.dtype, buffer=shm.buf)
+ #nan_corr_2d_mp function passes the frame and other details to nan_corr_2d for processing.
+ def nan_corr_2d_mp(i,obj, neighbor_box, min_neighbors, half_res_y, verbose):
+ obj_tmp_shared[i]=nan_corr_2d(obj, neighbor_box, min_neighbors, half_res_y,verbose,full_output=False)
+ # _nan_corr_2d unfolds the arguments
+ global _nan_corr_2d_mp
+ def _nan_corr_2d_mp(args):
+ nan_corr_2d_mp(*args)
+
+ context=multiprocessing.get_context('fork')
+ #processes argumnet is not provided to pool as no. of processes is
+ #inferred from os.cpu_count()
+ pool=context.Pool(processes=nproc, maxtasksperchild=1)
+ #creating a list of arguments
+ args=[]
+ for j in range(n_z):
+ args.append([j, obj_tmp[j], neighbor_box, min_neighbors, half_res_y, verbose])
+ try:
+ pool.map_async(_nan_corr_2d_mp,args, chunksize=1 ).get(timeout=10_000_000)
+ finally:
+ pool.close()
+ pool.join()
+ obj_tmp[:]=obj_tmp_shared[:]
+ shm.close()
+ shm.unlink()
+ #Original Multiprocessing code commented out below and ADACS approach is activated.
+ # res = pool_map(nproc, nan_corr_2d, iterable(obj_tmp), neighbor_box,
+ # min_neighbors, half_res_y, verbose, False)
+ # obj_tmp = np.array(res, dtype=object)
+
if verbose:
print('All nan pixels are corrected.')
