Skip to content

Commit

Permalink
Merge pull request #554 from VChristiaens/master
Browse files Browse the repository at this point in the history 
Addition of multiprocessing option for routines involving cube shifts
  • Loading branch information
@VChristiaens
VChristiaens authored Aug 25, 2022
2 parents fd7bd49 + 64d2a09 commit 20f5b51
Show file tree
Hide file tree
Showing 4 changed files with 139 additions and 90 deletions.
179 changes: 106 additions & 73 deletions vip_hci/fm/fakecomp.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -26,14 +27,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):
interpolation='lanczos4', transmission=None,
radial_gradient=False, full_output=False,
verbose=False, nproc=1):
""" Injects fake companions in branches, at given radial distances.
Parameters
Expand Down Expand Up @@ -93,6 +94,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
-------
Expand All @@ -105,24 +109,26 @@ 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)
"""
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,
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:
Expand Down Expand Up @@ -153,65 +159,36 @@ def _cube_inject_adi(array, psf_template, angle_list, flevel, plsc,

# 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')
-(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
Expand All @@ -230,11 +207,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:
Expand Down Expand Up @@ -345,6 +318,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.copy()


# 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)):
Expand Down
10 changes: 6 additions & 4 deletions vip_hci/metrics/contrcurve.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
-------
Expand All @@ -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)
Expand Down Expand Up @@ -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 +
Expand Down
34 changes: 23 additions & 11 deletions vip_hci/preproc/recentering.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
-------
Expand All @@ -283,15 +286,24 @@ 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)):
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 np.isscalar(shift_x):
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]):
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


Expand Down
6 changes: 4 additions & 2 deletions vip_hci/psfsub/pca_fullfr.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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]")
Expand All @@ -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")

0 comments on commit 20f5b51

Please sign in to comment.