From af64ed69d6d1bdab703de72de0a087869f8b8095 Mon Sep 17 00:00:00 2001 From: VChristiaens Date: Wed, 24 Aug 2022 16:41:51 +0200 Subject: [PATCH 1/7] Added the option of multiprocessing for cube shifts --- vip_hci/preproc/recentering.py | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/vip_hci/preproc/recentering.py b/vip_hci/preproc/recentering.py index b47818c5..8b6ee8ee 100644 --- a/vip_hci/preproc/recentering.py +++ b/vip_hci/preproc/recentering.py @@ -256,8 +256,8 @@ def frame_shift(array, shift_y, shift_x, imlib='vip-fft', return array_shifted -def cube_shift(cube, shift_y, shift_x, imlib='vip-fft', - interpolation='lanczos4', border_mode='reflect'): +def cube_shift(cube, shift_y, shift_x, imlib='vip-fft', + interpolation='lanczos4', border_mode='reflect', nproc=1): """ Shifts the X-Y coordinates of a cube or 3D array by x and y values. Parameters @@ -283,15 +283,21 @@ def cube_shift(cube, shift_y, shift_x, imlib='vip-fft', check_array(cube, dim=3) nfr = cube.shape[0] - cube_out = np.zeros_like(cube) - if isinstance(shift_x, (int, float)): + if np.isscalar(shift_x): shift_x = np.ones((nfr)) * shift_x if isinstance(shift_y, (int, float)): shift_y = np.ones((nfr)) * shift_y - for i in range(cube.shape[0]): - cube_out[i] = frame_shift(cube[i], shift_y[i], shift_x[i], imlib, - interpolation, border_mode) + if nproc == 1: + cube_out = np.zeros_like(cube) + for i in range(cube.shape[0]): + cube_out[i] = frame_shift(cube[i], shift_y[i], shift_x[i], imlib, + interpolation, border_mode) + elif nproc > 1: + res = pool_map(nproc, frame_shift, iterable(cube), iterable(shift_y), + iterable(shift_x), imlib, interpolation, border_mode) + cube_out = np.array(res) + return cube_out From f0319b4e7619c83b3cadff0b71393faddb4df31a Mon Sep 17 00:00:00 2001 From: VChristiaens Date: Wed, 24 Aug 2022 16:42:11 +0200 Subject: [PATCH 2/7] Added the option of multiprocessing for cube shifts within fake companion injection routine --- vip_hci/fm/fakecomp.py | 180 ++++++++++++++++++++++++----------------- 1 file changed, 105 insertions(+), 75 deletions(-) diff --git a/vip_hci/fm/fakecomp.py b/vip_hci/fm/fakecomp.py index 6edb7cbb..175946ed 100644 --- a/vip_hci/fm/fakecomp.py +++ b/vip_hci/fm/fakecomp.py @@ -26,14 +26,14 @@ frame_rotate) from ..var import (frame_center, fit_2dgaussian, fit_2dairydisk, fit_2dmoffat, get_circle, get_annulus_segments, dist_matrix) -from ..config.utils_conf import print_precision, check_array +from ..config.utils_conf import print_precision, check_array, pool_map, iterable def cube_inject_companions(array, psf_template, angle_list, flevel, rad_dists, - plsc=None, n_branches=1, theta=0, imlib='vip-fft', - interpolation='lanczos4', transmission=None, - radial_gradient=False, full_output=False, - verbose=False): + plsc=None, n_branches=1, theta=0, nproc=1, + imlib='vip-fft', interpolation='lanczos4', + transmission=None, radial_gradient=False, + full_output=False, verbose=False): """ Injects fake companions in branches, at given radial distances. Parameters @@ -72,6 +72,9 @@ def cube_inject_companions(array, psf_template, angle_list, flevel, rad_dists, Angle in degrees for rotating the position of the first branch that by default is located at zero degrees. Theta counts counterclockwise from the positive x axis. + nproc: int or None, optional + Number of CPUs to use for multiprocessing. If None, will be + automatically set to half the number of available CPUs. imlib : str, optional See the documentation of the ``vip_hci.preproc.frame_shift`` function. interpolation : str, optional @@ -105,24 +108,24 @@ def cube_inject_companions(array, psf_template, angle_list, flevel, rad_dists, [full_output & transmission != None] Array with injected psf affected by transmission (serves to check radial transmission) - """ + def _cube_inject_adi(array, psf_template, angle_list, flevel, plsc, rad_dists, n_branches=1, theta=0, imlib='vip-fft', interpolation='lanczos4', transmission=None, radial_gradient=False, verbose=False): - + if np.isscalar(flevel): + flevel = np.ones_like(angle_list)*flevel + if transmission is not None: # last radial separation should be beyond the edge of frame interp_trans = interp1d(transmission[0], transmission[1]) - positions = [] - w = int(np.ceil(size_fc/2)) - if size_fc % 2: # new convention - w -= 1 - sty = int(ceny) - w - stx = int(cenx) - w + ceny, cenx = frame_center(array[0]) + nframes = array.shape[-3] + size_fc = psf_template.shape[-1] + positions = [] # fake companion cube fc_fr = np.zeros([nframes, size_fc, size_fc]) if psf_template.ndim == 2: @@ -152,66 +155,37 @@ def _cube_inject_adi(array, psf_template, angle_list, flevel, plsc, fc_fr_rad[:, i] = interp_trans(d[i])*fc_fr[:, i] # check the effect of transmission on a single PSF tmp - psf_trans = frame_rotate(fc_fr_rad[0], - -(ang*180/np.pi-angle_list[0]), - imlib=imlib_rot, - interpolation=interpolation) - - shift_y = rad * np.sin(ang - np.deg2rad(angle_list[0])) - shift_x = rad * np.cos(ang - np.deg2rad(angle_list[0])) - dsy = shift_y-int(shift_y) - dsx = shift_x-int(shift_x) - fc_fr_ang = frame_shift(psf_trans, dsy, dsx, imlib_sh, - interpolation, - border_mode='constant') + # psf_trans = frame_rotate(fc_fr_rad[0], + # -(ang*180/np.pi-angle_list[0]), + # imlib=imlib_rot, + # interpolation=interpolation) + + # shift_y = rad * np.sin(ang - np.deg2rad(angle_list[0])) + # shift_x = rad * np.cos(ang - np.deg2rad(angle_list[0])) + # dsy = shift_y-int(shift_y) + # dsx = shift_x-int(shift_x) + # fc_fr_ang = frame_shift(psf_trans, dsy, dsx, imlib_sh, + # interpolation, + # border_mode='constant') else: fc_fr_rad = interp_trans(rad)*fc_fr - for fr in range(nframes): - shift_y = rad * np.sin(ang - np.deg2rad(angle_list[fr])) - shift_x = rad * np.cos(ang - np.deg2rad(angle_list[fr])) - if transmission is not None and radial_gradient: - fc_fr_ang = frame_rotate(fc_fr_rad[fr], - -(ang*180/np.pi - - angle_list[fr]), - imlib=imlib_rot, - interpolation=interpolation) - else: - fc_fr_ang = fc_fr_rad[fr] - - if np.isscalar(flevel): - fac = flevel - else: - fac = flevel[fr] - - # sub-px shift (within PSF template frame) - dsy = shift_y-int(shift_y) - dsx = shift_x-int(shift_x) - fc_fr_ang = frame_shift(fc_fr_ang, dsy, dsx, imlib_sh, - interpolation, - border_mode='constant') - # integer shift (in final cube) - y0 = sty+int(shift_y) - x0 = stx+int(shift_x) - yN = y0+size_fc - xN = x0+size_fc - p_y0 = 0 - p_x0 = 0 - p_yN = size_fc - p_xN = size_fc - if y0 < 0: - p_y0 = -y0 - y0 = 0 - if x0 < 0: - p_x0 = -x0 - x0 = 0 - if yN > sizey: - p_yN -= yN-sizey - yN = sizey - if xN > sizex: - p_xN -= xN-sizex - xN = sizex - array_out[fr, y0:yN, x0:xN] += fac*fc_fr_ang[p_y0:p_yN, - p_x0:p_xN] + if nproc == 1: + for fr in range(nframes): + array_out[fr] = _frame_shift_fcp(fc_fr_rad[fr], + array[fr], rad, ang, + angle_list[fr], + flevel[fr], size_fc, + imlib_sh, imlib_rot, + interpolation, + transmission, + radial_gradient) + else: + res = pool_map(nproc, _frame_shift_fcp, iterable(fc_fr_rad), + iterable(array), rad, ang, + iterable(angle_list), iterable(flevel), + size_fc, imlib_sh, imlib_rot, interpolation, + transmission, radial_gradient) + array_out = np.array(res) pos_y = rad * np.sin(ang) + ceny pos_x = rad * np.cos(ang) + cenx @@ -230,11 +204,7 @@ def _cube_inject_adi(array, psf_template, angle_list, flevel, plsc, check_array(psf_template, dim=(2, 3), msg="psf_template") nframes = array.shape[-3] - sizey = array.shape[-2] - sizex = array.shape[-1] - ceny, cenx = frame_center(array) - size_fc = psf_template.shape[-1] pceny, pcenx = frame_center(psf_template) if array.ndim == 4 and psf_template.ndim != 3: @@ -345,6 +315,66 @@ def _cube_inject_adi(array, psf_template, angle_list, flevel, plsc, return array_out +def _frame_shift_fcp(fc_fr_rad, array, rad, ang, derot_ang, flevel, size_fc, + imlib_sh, imlib_rot, interpolation, transmission, + radial_gradient): + """ + Specific cube shift algorithm for injection of fake companions + """ + + ceny, cenx = frame_center(array) + sizey = array.shape[-2] + sizex = array.shape[-1] + + array_sh = np.zeros_like(array) + + w = int(np.ceil(size_fc/2)) + if size_fc % 2: # new convention + w -= 1 + sty = int(ceny) - w + stx = int(cenx) - w + + shift_y = rad * np.sin(ang - np.deg2rad(derot_ang)) + shift_x = rad * np.cos(ang - np.deg2rad(derot_ang)) + if transmission is not None and radial_gradient: + fc_fr_ang = frame_rotate(fc_fr_rad, -(ang*180/np.pi -derot_ang), + imlib=imlib_rot, interpolation=interpolation) + else: + fc_fr_ang = fc_fr_rad + + + # sub-px shift (within PSF template frame) + dsy = shift_y-int(shift_y) + dsx = shift_x-int(shift_x) + fc_fr_ang = frame_shift(fc_fr_ang, dsy, dsx, imlib_sh, interpolation, + border_mode='constant') + # integer shift (in final cube) + y0 = sty+int(shift_y) + x0 = stx+int(shift_x) + yN = y0+size_fc + xN = x0+size_fc + p_y0 = 0 + p_x0 = 0 + p_yN = size_fc + p_xN = size_fc + if y0 < 0: + p_y0 = -y0 + y0 = 0 + if x0 < 0: + p_x0 = -x0 + x0 = 0 + if yN > sizey: + p_yN -= yN-sizey + yN = sizey + if xN > sizex: + p_xN -= xN-sizex + xN = sizex + + array_sh[y0:yN, x0:xN] += flevel*fc_fr_ang[p_y0:p_yN, p_x0:p_xN] + + return array_sh + + def generate_cube_copies_with_injections(array, psf_template, angle_list, plsc, n_copies=100, inrad=8, outrad=12, dist_flux=("uniform", 2, 500)): From addc104c86c9ed90025696e3f5e43173ec02260e Mon Sep 17 00:00:00 2001 From: VChristiaens Date: Wed, 24 Aug 2022 16:42:35 +0200 Subject: [PATCH 3/7] Added the option of multiprocessing for cube shifts within fake companion injection routine used in contrast curve calculation --- vip_hci/metrics/contrcurve.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/vip_hci/metrics/contrcurve.py b/vip_hci/metrics/contrcurve.py index b88980e0..5de27773 100644 --- a/vip_hci/metrics/contrcurve.py +++ b/vip_hci/metrics/contrcurve.py @@ -534,8 +534,8 @@ def throughput(cube, angle_list, psf_template, fwhm, algo, nbranch=1, theta=0, verbose : bool, optional If True prints out timing and information. **algo_dict - Parameters of the post-processing algorithms must be passed here, - including imlib and interpolation. + Parameters of the post-processing algorithms can be passed here, + including e.g. ``imlib``, ``interpolation`` or ``nproc``. Returns ------- @@ -562,6 +562,7 @@ def throughput(cube, angle_list, psf_template, fwhm, algo, nbranch=1, theta=0, """ array = cube parangles = angle_list + nproc = algo_dict.get('nproc', 1) imlib = algo_dict.get('imlib', 'vip-fft') interpolation = algo_dict.get('interpolation', 'lanczos4') scaling = algo_dict.get('scaling', None) @@ -706,8 +707,9 @@ def throughput(cube, angle_list, psf_template, fwhm, algo, nbranch=1, theta=0, rad_dists=[radvec[i]], theta=br*angle_branch + theta, - imlib=imlib, verbose=False, - interpolation=interpolation) + nproc=nproc, imlib=imlib, + interpolation=interpolation, + verbose=False) y = cy + radvec[i] * np.sin(np.deg2rad(br * angle_branch + theta)) x = cx + radvec[i] * np.cos(np.deg2rad(br * angle_branch + From 82bbe7b5131adde7f1a8f735c6fbf8b761b35340 Mon Sep 17 00:00:00 2001 From: VChristiaens Date: Wed, 24 Aug 2022 17:18:01 +0200 Subject: [PATCH 4/7] Better check of type of ncomp in project_subtract PCA function --- vip_hci/psfsub/pca_fullfr.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/vip_hci/psfsub/pca_fullfr.py b/vip_hci/psfsub/pca_fullfr.py index 1f2bf99d..ef809b8a 100644 --- a/vip_hci/psfsub/pca_fullfr.py +++ b/vip_hci/psfsub/pca_fullfr.py @@ -1071,7 +1071,7 @@ def _project_subtract(cube, cube_ref, ncomp, scaling, mask_center_px, svd_mode, vectors of the input matrix, as returned by ``svd/svd_wrapper()`` """ _, y, x = cube.shape - if isinstance(ncomp, int): + if isinstance(ncomp, (int, np.integer)): # if cube_sig is not None: # cube_emp = cube-cube_sig # else: @@ -1124,7 +1124,7 @@ def _project_subtract(cube, cube_ref, ncomp, scaling, mask_center_px, svd_mode, else: return residuals_res - elif isinstance(ncomp, float): + elif isinstance(ncomp, (float, np.float)): if not 1 > ncomp > 0: raise ValueError("when `ncomp` is float, it must lie in the " "interval (0,1]") @@ -1148,3 +1148,5 @@ def _project_subtract(cube, cube_ref, ncomp, scaling, mask_center_px, svd_mode, return residuals_res, reconstructed, V else: return residuals_res + else: + raise TypeError("Type not recognized for ncomp, should be int or float") \ No newline at end of file From 106dd96bdb8ce52959314aff324749afe7b4c88f Mon Sep 17 00:00:00 2001 From: VChristiaens Date: Wed, 24 Aug 2022 17:49:24 +0200 Subject: [PATCH 5/7] Changed location of new parameter nproc in cube_inject_companions for retro-compatibility --- vip_hci/fm/fakecomp.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/vip_hci/fm/fakecomp.py b/vip_hci/fm/fakecomp.py index 175946ed..c7f0eaec 100644 --- a/vip_hci/fm/fakecomp.py +++ b/vip_hci/fm/fakecomp.py @@ -30,10 +30,10 @@ def cube_inject_companions(array, psf_template, angle_list, flevel, rad_dists, - plsc=None, n_branches=1, theta=0, nproc=1, - imlib='vip-fft', interpolation='lanczos4', - transmission=None, radial_gradient=False, - full_output=False, verbose=False): + plsc=None, n_branches=1, theta=0, imlib='vip-fft', + interpolation='lanczos4', transmission=None, + radial_gradient=False, full_output=False, + verbose=False, nproc=1): """ Injects fake companions in branches, at given radial distances. Parameters @@ -72,9 +72,6 @@ def cube_inject_companions(array, psf_template, angle_list, flevel, rad_dists, Angle in degrees for rotating the position of the first branch that by default is located at zero degrees. Theta counts counterclockwise from the positive x axis. - nproc: int or None, optional - Number of CPUs to use for multiprocessing. If None, will be - automatically set to half the number of available CPUs. imlib : str, optional See the documentation of the ``vip_hci.preproc.frame_shift`` function. interpolation : str, optional @@ -96,6 +93,9 @@ def cube_inject_companions(array, psf_template, angle_list, flevel, rad_dists, to the new array. verbose : bool, optional If True prints out additional information. + nproc: int or None, optional + Number of CPUs to use for multiprocessing. If None, will be + automatically set to half the number of available CPUs. Returns ------- From c7e411c2fdab91f0bdb4ef6c2c9a7f7d5a608085 Mon Sep 17 00:00:00 2001 From: VChristiaens Date: Wed, 24 Aug 2022 22:19:43 +0200 Subject: [PATCH 6/7] implementation of multiprocessing-based fake companion injection --- vip_hci/fm/fakecomp.py | 35 +++++++++++++++++++---------------- 1 file changed, 19 insertions(+), 16 deletions(-) diff --git a/vip_hci/fm/fakecomp.py b/vip_hci/fm/fakecomp.py index c7f0eaec..a1d9ab9c 100644 --- a/vip_hci/fm/fakecomp.py +++ b/vip_hci/fm/fakecomp.py @@ -11,6 +11,7 @@ 'generate_cube_copies_with_injections', 'frame_inject_companion'] +from multiprocessing import cpu_count import numpy as np from scipy import stats from scipy.interpolate import interp1d @@ -109,7 +110,9 @@ def cube_inject_companions(array, psf_template, angle_list, flevel, rad_dists, by transmission (serves to check radial transmission) """ - + if nproc is None: + nproc = cpu_count()//2 + def _cube_inject_adi(array, psf_template, angle_list, flevel, plsc, rad_dists, n_branches=1, theta=0, imlib='vip-fft', interpolation='lanczos4', transmission=None, @@ -155,10 +158,10 @@ def _cube_inject_adi(array, psf_template, angle_list, flevel, plsc, fc_fr_rad[:, i] = interp_trans(d[i])*fc_fr[:, i] # check the effect of transmission on a single PSF tmp - # psf_trans = frame_rotate(fc_fr_rad[0], - # -(ang*180/np.pi-angle_list[0]), - # imlib=imlib_rot, - # interpolation=interpolation) + psf_trans = frame_rotate(fc_fr_rad[0], + -(ang*180/np.pi-angle_list[0]), + imlib=imlib_rot, + interpolation=interpolation) # shift_y = rad * np.sin(ang - np.deg2rad(angle_list[0])) # shift_x = rad * np.cos(ang - np.deg2rad(angle_list[0])) @@ -171,21 +174,21 @@ def _cube_inject_adi(array, psf_template, angle_list, flevel, plsc, fc_fr_rad = interp_trans(rad)*fc_fr if nproc == 1: for fr in range(nframes): - array_out[fr] = _frame_shift_fcp(fc_fr_rad[fr], - array[fr], rad, ang, - angle_list[fr], - flevel[fr], size_fc, - imlib_sh, imlib_rot, - interpolation, - transmission, - radial_gradient) + array_out[fr] += _frame_shift_fcp(fc_fr_rad[fr], + array[fr], rad, ang, + angle_list[fr], + flevel[fr], size_fc, + imlib_sh, imlib_rot, + interpolation, + transmission, + radial_gradient) else: res = pool_map(nproc, _frame_shift_fcp, iterable(fc_fr_rad), iterable(array), rad, ang, iterable(angle_list), iterable(flevel), size_fc, imlib_sh, imlib_rot, interpolation, transmission, radial_gradient) - array_out = np.array(res) + array_out += np.array(res) pos_y = rad * np.sin(ang) + ceny pos_x = rad * np.cos(ang) + cenx @@ -340,7 +343,7 @@ def _frame_shift_fcp(fc_fr_rad, array, rad, ang, derot_ang, flevel, size_fc, fc_fr_ang = frame_rotate(fc_fr_rad, -(ang*180/np.pi -derot_ang), imlib=imlib_rot, interpolation=interpolation) else: - fc_fr_ang = fc_fr_rad + fc_fr_ang = fc_fr_rad.copy() # sub-px shift (within PSF template frame) @@ -370,7 +373,7 @@ def _frame_shift_fcp(fc_fr_rad, array, rad, ang, derot_ang, flevel, size_fc, p_xN -= xN-sizex xN = sizex - array_sh[y0:yN, x0:xN] += flevel*fc_fr_ang[p_y0:p_yN, p_x0:p_xN] + array_sh[y0:yN, x0:xN] = flevel*fc_fr_ang[p_y0:p_yN, p_x0:p_xN] return array_sh From 64d2a09ce28db827270c2e1230f07ffc1516a098 Mon Sep 17 00:00:00 2001 From: VChristiaens Date: Wed, 24 Aug 2022 22:20:15 +0200 Subject: [PATCH 7/7] Improved description for cube_shift --- vip_hci/preproc/recentering.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/vip_hci/preproc/recentering.py b/vip_hci/preproc/recentering.py index 8b6ee8ee..2b5a2096 100644 --- a/vip_hci/preproc/recentering.py +++ b/vip_hci/preproc/recentering.py @@ -273,6 +273,9 @@ def cube_shift(cube, shift_y, shift_x, imlib='vip-fft', See the documentation of the ``vip_hci.preproc.frame_shift`` function. border_mode : str, optional See the documentation of the ``vip_hci.preproc.frame_shift`` function. + nproc: int or None, optional + Number of CPUs to use for multiprocessing. If None, will be + automatically set to half the number of available CPUs. Returns ------- @@ -284,10 +287,13 @@ def cube_shift(cube, shift_y, shift_x, imlib='vip-fft', nfr = cube.shape[0] if np.isscalar(shift_x): - shift_x = np.ones((nfr)) * shift_x - if isinstance(shift_y, (int, float)): - shift_y = np.ones((nfr)) * shift_y + shift_x = np.ones([nfr]) * shift_x + if np.isscalar(shift_y): + shift_y = np.ones([nfr]) * shift_y + if nproc is None: + nproc = cpu_count()//2 + if nproc == 1: cube_out = np.zeros_like(cube) for i in range(cube.shape[0]):