From 022bde0da25ddfbb054043bcd3a85ef25edc62fe Mon Sep 17 00:00:00 2001 From: Carlos Gomez Date: Tue, 13 Feb 2018 14:33:15 +0100 Subject: [PATCH] VIP v0.8.8. Fixed multiprocessing for annular PCA, LLSG and ``snrmap`` functions in Python 3.6. Added support for ADI+IFS data to ``throughput``, ``contrast_curve`` and ``psf_norm`` functions. Fixed full-frame PCA for ADI+IFS data. Fixed ``frame_center`` function for even-sized frames. Fixed frame and cube cropping for even-sized frames. ``cube_crop_frames`` now supports ADI+IFS data. --- docs/source/faq.rst | 20 +- docs/source/readme.rst | 2 +- readme.rst | 2 +- vip_hci/__init__.py | 2 +- vip_hci/llsg/llsg.py | 34 +-- vip_hci/negfc/mcmc_sampling.py | 4 +- vip_hci/negfc/simplex_fmerit.py | 4 +- vip_hci/negfc/simplex_optim.py | 4 +- vip_hci/pca/pca_fullfr.py | 30 +-- vip_hci/pca/pca_local.py | 68 ++--- vip_hci/pca/utils_pca.py | 15 +- vip_hci/phot/contrcurve.py | 423 +++++++++++++++++++++---------- vip_hci/phot/fakecomp.py | 198 +++++++++------ vip_hci/phot/snr.py | 95 ++++--- vip_hci/preproc/__init__.py | 1 - vip_hci/preproc/cosmetics.py | 306 ++++++++++++++++++---- vip_hci/preproc/cosmetics_ifs.py | 201 --------------- vip_hci/preproc/derotation.py | 9 +- vip_hci/preproc/recentering.py | 59 +++-- vip_hci/var/shapes.py | 47 ++-- vip_hci/var/utils_var.py | 142 +++++++---- 21 files changed, 950 insertions(+), 716 deletions(-) delete mode 100644 vip_hci/preproc/cosmetics_ifs.py diff --git a/docs/source/faq.rst b/docs/source/faq.rst index 4ce6c52f..050b8c6e 100644 --- a/docs/source/faq.rst +++ b/docs/source/faq.rst @@ -23,21 +23,11 @@ matplotlibrc file is located in: $HOME/.config/matplotlib/matplotlibrc -* Why do I get, in OSX, the following RuntimeError? -*Python is not installed as a framework. The Mac OS X backend will not be able -to function correctly if Python is not installed as a framework. See the -Python documentation for more information on installing Python as a framework -on Mac OS X. Please either reinstall Python as a framework, or try one of the -other backends.* -Again, this is a matplotlib-backend related issue. Read the link in the previous -answer. It can be solved setting the backend to WXAgg or TkAgg. Optionally, you -can change your matplotlib backend **before** importing ``VIP``: - -.. code-block:: python - - import matplotlib as mpl - mpl.use('TkAgg') - import vip +* I get a RuntimeError (Python is not installed as a framework) in Macosx, +when using python3 and conda. Why? +This is caused by using matplotlib with virtual environments or conda in Macosx, +and has nothing to do with VIP. See 'Working with Matplotlib on OSX' in the +Matplotlib FAQ for more information. * The ``VIP`` setup.py script doesn't finish the job, it seems to be stuck. What diff --git a/docs/source/readme.rst b/docs/source/readme.rst index 194d0e68..de8686a8 100644 --- a/docs/source/readme.rst +++ b/docs/source/readme.rst @@ -42,7 +42,7 @@ made by collaborators from several teams (take a look at the tab contributors on it doesn't mean it's free from bugs. The code is continuously evolving and therefore feedback/contributions are greatly appreciated. If you want to report a bug, suggest or add a functionality please create an issue or send a pull -request on `this repository `_. +request `here `_. Documentation diff --git a/readme.rst b/readme.rst index 194d0e68..de8686a8 100644 --- a/readme.rst +++ b/readme.rst @@ -42,7 +42,7 @@ made by collaborators from several teams (take a look at the tab contributors on it doesn't mean it's free from bugs. The code is continuously evolving and therefore feedback/contributions are greatly appreciated. If you want to report a bug, suggest or add a functionality please create an issue or send a pull -request on `this repository `_. +request `here `_. Documentation diff --git a/vip_hci/__init__.py b/vip_hci/__init__.py index dfb37487..bdb4824a 100644 --- a/vip_hci/__init__.py +++ b/vip_hci/__init__.py @@ -12,7 +12,7 @@ from . import stats from . import var -__version__ = "0.8.7" +__version__ = "0.8.8" print("---------------------------------------------------") diff --git a/vip_hci/llsg/llsg.py b/vip_hci/llsg/llsg.py index 3b256b0a..086106b5 100644 --- a/vip_hci/llsg/llsg.py +++ b/vip_hci/llsg/llsg.py @@ -31,7 +31,7 @@ def llsg(cube, angle_list, fwhm, rank=10, thresh=1, max_iter=10, residuals_tol=1e-1, cevr=0.9, thresh_mode='soft', nproc=1, asize=None, n_segments=4, azimuth_overlap=None, radius_int=None, random_seed=None, imlib='opencv', interpolation='lanczos4', - high_pass=False, collapse='median', full_output=False, verbose=True, + high_pass=None, collapse='median', full_output=False, verbose=True, debug=False): """ Local Low-rank plus Sparse plus Gaussian-noise decomposition (LLSG) as described in Gomez Gonzalez et al. 2016. This first version of our algorithm @@ -228,22 +228,22 @@ def llsg(cube, angle_list, fwhm, rank=10, thresh=1, max_iter=10, else: pool = Pool(processes=int(nproc)) - res = pool.map(EFT, itt.izip(itt.repeat(_decompose_patch), - itt.repeat(indices), - range(n_segments_ann), - itt.repeat(n_segments_ann), - itt.repeat(rank), - itt.repeat(low_rank_ref), - itt.repeat(low_rank_mode), - itt.repeat(thresh), - itt.repeat(thresh_mode), - itt.repeat(max_iter), - itt.repeat(auto_rank_mode), - itt.repeat(cevr), - itt.repeat(residuals_tol), - itt.repeat(random_seed), - itt.repeat(debug), - itt.repeat(full_output))) + res = pool.map(EFT, zip(itt.repeat(_decompose_patch), + itt.repeat(indices), + range(n_segments_ann), + itt.repeat(n_segments_ann), + itt.repeat(rank), + itt.repeat(low_rank_ref), + itt.repeat(low_rank_mode), + itt.repeat(thresh), + itt.repeat(thresh_mode), + itt.repeat(max_iter), + itt.repeat(auto_rank_mode), + itt.repeat(cevr), + itt.repeat(residuals_tol), + itt.repeat(random_seed), + itt.repeat(debug), + itt.repeat(full_output))) patches = np.array(res) pool.close() diff --git a/vip_hci/negfc/mcmc_sampling.py b/vip_hci/negfc/mcmc_sampling.py index c3858100..4950e4c3 100644 --- a/vip_hci/negfc/mcmc_sampling.py +++ b/vip_hci/negfc/mcmc_sampling.py @@ -6,9 +6,7 @@ from __future__ import print_function __author__ = 'O. Wertz, Carlos Alberto Gomez Gonzalez' -__all__ = ['lnprior', - 'lnlike', - 'mcmc_negfc_sampling', +__all__ = ['mcmc_negfc_sampling', 'chain_zero_truncated', 'show_corner_plot', 'show_walk_plot', diff --git a/vip_hci/negfc/simplex_fmerit.py b/vip_hci/negfc/simplex_fmerit.py index d260e1f3..2136eb7b 100644 --- a/vip_hci/negfc/simplex_fmerit.py +++ b/vip_hci/negfc/simplex_fmerit.py @@ -5,11 +5,13 @@ """ from __future__ import print_function +__all__ = [] + import numpy as np from skimage.draw import circle from ..phot import cube_inject_companions from ..var import frame_center -from ..pca import pca_annulus +from ..pca.utils_pca import pca_annulus def chisquare(modelParameters, cube, angs, plsc, psfs_norm, fwhm, annulus_width, diff --git a/vip_hci/negfc/simplex_optim.py b/vip_hci/negfc/simplex_optim.py index 71aa2e69..95d1f7d5 100644 --- a/vip_hci/negfc/simplex_optim.py +++ b/vip_hci/negfc/simplex_optim.py @@ -13,9 +13,7 @@ from ..var import frame_center from ..conf import time_ini, timing, sep -__all__ = ['firstguess_from_coord', - 'firstguess_simplex', - 'firstguess'] +__all__ = ['firstguess'] diff --git a/vip_hci/pca/pca_fullfr.py b/vip_hci/pca/pca_fullfr.py index a5c0a2f3..bfc7b077 100644 --- a/vip_hci/pca/pca_fullfr.py +++ b/vip_hci/pca/pca_fullfr.py @@ -96,7 +96,8 @@ def pca(cube, angle_list=None, cube_ref=None, scale_list=None, ncomp=1, ncomp2=1 ncomp is the number of PCs from the library of spectral channels. ncomp2 : int, optional How many PCs are used for IFS+ADI datacubes in the second stage PCA. - ncomp2 goes up to the number of multi-spectral frames. + ncomp2 goes up to the number ADI frames. If None then the second stage + of the PCA is ignored and the residuals are derotated and combined. svd_mode : {'lapack', 'arpack', 'eigen', 'randsvd', 'cupy', 'eigencupy', 'randcupy'}, str Switch for the SVD method/library to be used. ``lapack`` uses the LAPACK linear algebra library through Numpy and it is the most conventional way @@ -193,7 +194,7 @@ def subtract_projection(cube, cube_ref, ncomp, scaling, mask_center_px, else: ref_lib = matrix - if indices is not None and frame is not None: # one row (frame) at a time + if indices is not None and frame is not None: # one row (frame) at a time ref_lib = ref_lib[indices] if ref_lib.shape[0] <= 10: msg = 'Too few frames left in the PCA library (<10). ' @@ -209,7 +210,7 @@ def subtract_projection(cube, cube_ref, ncomp, scaling, mask_center_px, return ref_lib.shape[0], residuals, reconstructed else: return ref_lib.shape[0], residuals - else: # the whole matrix + else: # the whole matrix V = svd_wrapper(ref_lib, svd_mode, ncomp, debug, verbose) if verbose: timing(start_time) @@ -263,7 +264,8 @@ def subtract_projection(cube, cube_ref, ncomp, scaling, mask_center_px, if not scale_list.shape[0]==cube.shape[0]: raise TypeError('Scaling factors vector has wrong length') - if verbose: start_time = time_ini() + if verbose: + start_time = time_ini() if check_mem: input_bytes = cube.nbytes @@ -322,7 +324,7 @@ def subtract_projection(cube, cube_ref, ncomp, scaling, mask_center_px, # SDI (IFS) + ADI: cube with multiple spectral channels + rotation # shape of cube: [# channels, # adi-frames, Y, X] #*********************************************************************** - elif cube.ndim==4 and angle_list is not None: + elif cube.ndim == 4 and angle_list is not None: if verbose: print('{:} spectral channels in IFS cube'.format(z)) print('{:} ADI frames in all channels'.format(n)) @@ -353,22 +355,20 @@ def subtract_projection(cube, cube_ref, ncomp, scaling, mask_center_px, residuals_cube_channels[i] = frame bar.update() - # de-rotation of the PCA processed channels - if ncomp2 > z: - ncomp2 = min(ncomp2, z) - msg = 'Number of PCs too high (max PCs={}), using instead {:} PCs.' - print(msg.format(n, ncomp2)) - - elif ncomp2 is None: + # de-rotation of the PCA processed channels, ADI fashion + if ncomp2 is None: residuals_cube_channels_ = cube_derotate(residuals_cube_channels, - angle_list, imlib=imlib, - interpolation=interpolation) + angle_list, imlib=imlib, + interpolation=interpolation) frame = cube_collapse(residuals_cube_channels_, mode=collapse) if verbose: print('De-rotating and combining') timing(start_time) - else: + if ncomp2 > n: + ncomp2 = min(ncomp2, n) + msg = 'Number of PCs too high (max PCs={}), using instead {:} PCs.' + print(msg.format(n, ncomp2)) res_ifs_adi = subtract_projection(residuals_cube_channels, None, ncomp2, scaling, mask_center_px, debug, svd_mode, False, diff --git a/vip_hci/pca/pca_local.py b/vip_hci/pca/pca_local.py index 930bcb34..08636692 100644 --- a/vip_hci/pca/pca_local.py +++ b/vip_hci/pca/pca_local.py @@ -22,12 +22,11 @@ from ..stats import descriptive_stats - def pca_rdi_annular(cube, angle_list, cube_ref, radius_int=0, asize=1, ncomp=1, svd_mode='randsvd', min_corr=0.9, fwhm=4, scaling='temp-standard', imlib='opencv', interpolation='lanczos4', collapse='median', - full_output=False, verbose=True, debug=False): + full_output=False, verbose=True): """ Annular PCA with Reference Library + Correlation + standardization In the case of having a large number of reference images, e.g. for a survey @@ -87,9 +86,7 @@ def pca_rdi_annular(cube, angle_list, cube_ref, radius_int=0, asize=1, Whether to return the final median combined image only or with other intermediate arrays. verbose : {True, False}, bool optional - If True prints to stdout intermediate info. - debug : {False, True}, bool optional - Whether to output some intermediate information. + If True prints to stdout intermediate info. Returns ------- @@ -204,7 +201,7 @@ def do_pca_annulus(ncomp, matrix, svd_mode, noise_error, data_ref): return frame -def pca_adi_annular(cube, angle_list, radius_int=0, fwhm=4, asize=3, +def pca_adi_annular(cube, angle_list, radius_int=0, fwhm=4, asize=3, delta_rot=1, ncomp=1, svd_mode='randsvd', nproc=1, min_frames_pca=10, tol=1e-1, scaling=None, quad=False, imlib='opencv', interpolation='lanczos4', collapse='median', @@ -219,10 +216,7 @@ def pca_adi_annular(cube, angle_list, radius_int=0, fwhm=4, asize=3, algebra calcularions, which comes by default in every OSX system, is broken for multiprocessing. Avoid using this function unless you have compiled Python against other linear algebra library. An easy fix is to install - latest ANACONDA (2.5 or later) distribution which ships MKL library - (replacing the problematic ACCELERATE). On linux with the default - LAPACK/BLAS libraries it successfully distributes the processes among all - the existing cores. + (ana)conda and the openblas or MKL libraries (replacing ACCELERATE). Parameters ---------- @@ -378,10 +372,10 @@ def pca_adi_annular(cube, angle_list, radius_int=0, fwhm=4, asize=3, #*************************************************************** # We loop the frames and do the PCA to obtain the residuals cube #*************************************************************** - residuals = do_pca_loop(matrix_quad, yy, xx, nproc, angle_list, - fwhm, pa_thr, scaling, ann_center, - svd_mode, ncompann, min_frames_pca, tol, - debug, verbose) + residuals = do_pca_loop(matrix_quad, nproc, angle_list, fwhm, + pa_thr, scaling, ann_center, svd_mode, + ncompann, min_frames_pca, tol, debug, + verbose) for frame in range(n): cube_out[frame][yy, xx] = residuals[frame] @@ -404,9 +398,9 @@ def pca_adi_annular(cube, angle_list, radius_int=0, fwhm=4, asize=3, #******************************************************************* # We loop the frames and do the PCA to obtain the residuals cube #******************************************************************* - residuals = do_pca_loop(matrix_ann, yy, xx, nproc, angle_list, fwhm, - pa_thr, scaling, ann_center, svd_mode, - ncompann, min_frames_pca, tol, debug, verbose) + residuals = do_pca_loop(matrix_ann, nproc, angle_list, fwhm, pa_thr, + scaling, ann_center, svd_mode, ncompann, + min_frames_pca, tol, debug, verbose) for frame in range(n): cube_out[frame][yy, xx] = residuals[frame] @@ -504,11 +498,11 @@ def find_indices(angle_list, frame, thr, truncate): else: half2 = range(index_foll, min(n, int(thr/2+index_foll))) half1 = range(max(0,index_prev-thr+len(half2)), index_prev) - return np.array(half1+half2) + return np.array(list(half1) + list(half2)) -def do_pca_loop(matrix, yy, xx, nproc, angle_list, fwhm, pa_threshold, scaling, +def do_pca_loop(matrix, nproc, angle_list, fwhm, pa_threshold, scaling, ann_center, svd_mode, ncomp, min_frames_pca, tol, debug, verbose): """ """ @@ -521,40 +515,34 @@ def do_pca_loop(matrix, yy, xx, nproc, angle_list, fwhm, pa_threshold, scaling, #*************************************************************** ncomps = [] nfrslib = [] - if nproc==1: + if nproc == 1: residualarr = [] for frame in range(n): res = do_pca_patch(matrix_ann, frame, angle_list, fwhm, - pa_threshold, scaling, ann_center, - svd_mode, ncomp, min_frames_pca, tol, - debug) + pa_threshold, ann_center, svd_mode, ncomp, + min_frames_pca, tol, debug) residualarr.append(res[0]) ncomps.append(res[1]) nfrslib.append(res[2]) residuals = np.array(residualarr) - elif nproc>1: + elif nproc > 1: #*********************************************************** # A multiprocessing pool is created to process the frames in # a parallel way. SVD/PCA is done in do_pca_patch function #*********************************************************** pool = Pool(processes=int(nproc)) - res = pool.map(EFT, itt.izip(itt.repeat(do_pca_patch), - itt.repeat(matrix_ann), - range(n), itt.repeat(angle_list), - itt.repeat(fwhm), - itt.repeat(pa_threshold), - itt.repeat(scaling), - itt.repeat(ann_center), - itt.repeat(svd_mode), - itt.repeat(ncomp), - itt.repeat(min_frames_pca), - itt.repeat(tol), - itt.repeat(debug))) + res = pool.map(EFT, zip(itt.repeat(do_pca_patch), + itt.repeat(matrix_ann), range(n), + itt.repeat(angle_list), itt.repeat(fwhm), + itt.repeat(pa_threshold), + itt.repeat(ann_center), itt.repeat(svd_mode), + itt.repeat(ncomp), itt.repeat(min_frames_pca), + itt.repeat(tol))) res = np.array(res) residuals = np.array(res[:,0]) - ncomps = res[:,1] - nfrslib = res[:,2] + ncomps = res[:, 1] + nfrslib = res[:, 2] pool.close() # number of frames in library printed for each annular quadrant @@ -567,8 +555,8 @@ def do_pca_loop(matrix, yy, xx, nproc, angle_list, fwhm, pa_threshold, scaling, return residuals -def do_pca_patch(matrix, frame, angle_list, fwhm, pa_threshold, scaling, - ann_center, svd_mode, ncomp, min_frames_pca, tol, debug): +def do_pca_patch(matrix, frame, angle_list, fwhm, pa_threshold, ann_center, + svd_mode, ncomp, min_frames_pca, tol): """ Does the SVD/PCA for each frame patch (small matrix). For each frame we find the frames to be rejected depending on the amount of rotation. The diff --git a/vip_hci/pca/utils_pca.py b/vip_hci/pca/utils_pca.py index be27a211..776b8d20 100644 --- a/vip_hci/pca/utils_pca.py +++ b/vip_hci/pca/utils_pca.py @@ -8,8 +8,7 @@ from __future__ import print_function __author__ = 'Carlos Alberto Gomez Gonzalez' -__all__ = ['pca_annulus', - 'scale_cube_for_pca'] +__all__ = [] import numpy as np from ..var import get_square_robust, frame_center, prepare_matrix @@ -42,9 +41,8 @@ def scale_cube_for_pca(cube,scal_list, full_output=True, inverse=False, y_in=1, or not; i.e. True is to descale the cube (typically after a first scaling has already been done) y_in, x-in: - Initial y and x sizes. - In case the cube is descaled, these values will be used to crop back the - cubes/frames to their original size. + Initial y and x sizes. In case the cube is descaled, these values will + be used to crop back the cubes/frames to their original size. imlib : str optional See the documentation of the ``vip_hci.preproc.frame_rescaling`` function. @@ -92,21 +90,22 @@ def scale_cube_for_pca(cube,scal_list, full_output=True, inverse=False, y_in=1, cy,cx = frame_center(cube[0]) # (de)scale the cube, so that a planet would now move radially - cube,frame = cube_rescaling(big_cube,var_list,ref_y=cy, ref_x=cx, - imlib=imlib, interpolation=interpolation) + cube, frame = cube_rescaling(big_cube, var_list, ref_y=cy, ref_x=cx, + imlib=imlib, interpolation=interpolation) if inverse: if max_sc > 1: + # TODO: check the use of get_square_robust frame = get_square_robust(frame,max(y_in,x_in), cy,cx,strict=False) if full_output: n_z = cube.shape[0] array_old = cube.copy() cube = np.zeros([n_z,max(y_in,x_in),max(y_in,x_in)]) for zz in range(n_z): + # TODO: check the use of get_square_robust cube[zz]=get_square_robust(array_old[zz],max(y_in,x_in), cy,cx,strict=False) - if full_output: return cube,frame,y,x,cy,cx else: diff --git a/vip_hci/phot/contrcurve.py b/vip_hci/phot/contrcurve.py index 47fc61ab..879eaf9f 100644 --- a/vip_hci/phot/contrcurve.py +++ b/vip_hci/phot/contrcurve.py @@ -26,29 +26,31 @@ from ..var import frame_center, dist - def contrast_curve(cube, angle_list, psf_template, fwhm, pxscale, starphot, - algo, sigma=5, nbranch=1, theta=0, inner_rad=1, wedge=(0,360), - fc_snr=10.0, student=True, transmission=None, smooth=True, - plot=True, dpi=100, imlib='opencv', debug=False, verbose=True, full_output=False, - save_plot=None, object_name=None, frame_size=None, - fix_y_lim=(), figsize=(8,4), **algo_dict): - """ Computes the contrast curve for a given SIGMA (*sigma*) level. The - contrast is calculated as sigma*noise/throughput. This implementation takes - into account the small sample statistics correction proposed in Mawet et al. - 2014. + algo, sigma=5, nbranch=1, theta=0, inner_rad=1, + wedge=(0,360), fc_snr=10.0, student=True, transmission=None, + smooth=True, plot=True, dpi=100, imlib='opencv', debug=False, + verbose=True, full_output=False, save_plot=None, + object_name=None, frame_size=None, fix_y_lim=(), + figsize=(8,4), **algo_dict): + """ Computes the contrast curve at a given SIGMA (``sigma``) level for an + ADI cube or ADI+IFS cube. The contrast is calculated as + sigma*noise/throughput. This implementation takes into account the small + sample statistics correction proposed in Mawet et al. 2014. Parameters ---------- cube : array_like - The input cube without fake companions. + The input cube, 2d (ADI data) or 3d array (IFS data), without fake + companions. angle_list : array_like Vector with the parallactic angles. psf_template : array_like Frame with the psf template for the fake companion(s). PSF must be centered in array. Normalization is done internally. - fwhm : float - FWHM in pixels. + fwhm: int or float or 1d array, optional + The the Full Width Half Maximum in pixels. It can handle a different + FWHM value for different wavelengths (IFS data). pxscale : float Plate scale or pixel scale of the instrument. starphot : int or float or 1d array @@ -133,16 +135,21 @@ def contrast_curve(cube, angle_list, psf_template, fwhm, pxscale, starphot, 3d array with 3 frames containing the position of the companions in the 3 patterns. """ - if not cube.ndim == 3: - raise TypeError('The input array is not a cube') - if not cube.shape[0] == angle_list.shape[0]: - raise TypeError('Input vector or parallactic angles has wrong length') - if not psf_template.ndim==2: - raise TypeError('Template PSF is not a frame') + if not (cube.ndim == 3 or cube.ndim == 4): + raise TypeError('The input array is not a 3d or 4d cube') + if cube.ndim == 3 and (not cube.shape[0] == angle_list.shape[0]): + raise TypeError('Input parallactic angles vector has wrong length') + if cube.ndim == 4 and (not cube.shape[1] == angle_list.shape[0]): + raise TypeError('Input parallactic angles vector has wrong length') + if cube.ndim == 3 and not psf_template.ndim == 2: + raise TypeError('Template PSF is not a frame (for ADI case)') + if cube.ndim == 4 and not psf_template.ndim == 3: + raise TypeError('Template PSF is not a cube (for ADI+IFS case)') if transmission is not None: - if not isinstance(transmission, tuple) or not len(transmission)==2: + if not isinstance(transmission, tuple) or not len(transmission) == 2: raise TypeError('transmission must be a tuple with 2 1d vectors') - if isinstance(starphot, float) or isinstance(starphot, int): pass + if isinstance(starphot, float) or isinstance(starphot, int): + pass else: if not starphot.shape[0] == cube.shape[0]: raise TypeError('Correction vector has bad size') @@ -163,14 +170,18 @@ def contrast_curve(cube, angle_list, psf_template, fwhm, pxscale, starphot, # throughput verbose_thru = False - if verbose==2: verbose_thru = True + if verbose == 2: + verbose_thru = True res_throug = throughput(cube, angle_list, psf_template, fwhm, pxscale, nbranch=nbranch, theta=theta, inner_rad=inner_rad, - wedge=wedge, fc_snr=fc_snr, full_output=True, algo=algo, - imlib=imlib, verbose=verbose_thru, **algo_dict) + wedge=wedge, fc_snr=fc_snr, full_output=True, + algo=algo, imlib=imlib, verbose=verbose_thru, + **algo_dict) vector_radd = res_throug[2] - if res_throug[0].shape[0]>1: thruput_mean = np.mean(res_throug[0], axis=0) - else: thruput_mean = res_throug[0][0] + if res_throug[0].shape[0] > 1: + thruput_mean = np.mean(res_throug[0], axis=0) + else: + thruput_mean = res_throug[0][0] cube_fc_all = res_throug[3] frame_fc_all = res_throug[4] frame_nofc = res_throug[5] @@ -180,7 +191,7 @@ def contrast_curve(cube, angle_list, psf_template, fwhm, pxscale, starphot, print('Finished the throughput calculation') timing(start_time) - if thruput_mean[-1]==0: + if thruput_mean[-1] == 0: thruput_mean = thruput_mean[:-1] vector_radd = vector_radd[:-1] @@ -210,7 +221,8 @@ def contrast_curve(cube, angle_list, psf_template, fwhm, pxscale, starphot, if smooth: # smoothing the noise vector using a Savitzky-Golay filter win = min(noise_samp.shape[0]-2,int(2*fwhm)) - if win%2==0.: win += 1 + if win%2 == 0.: + win += 1 noise_samp_sm = savgol_filter(noise_samp, polyorder=2, mode='nearest', window_length=win) else: @@ -297,7 +309,8 @@ def contrast_curve(cube, angle_list, psf_template, fwhm, pxscale, starphot, pca_type = 'ADI' else: pca_type = 'RDI' - title = pca_type + ' ' + object_name + ' ' + str(ncomp) + 'pc ' + str(frame_size) + '+' + str(inner_rad) + title = pca_type + ' ' + object_name + ' ' + str(ncomp) + title += 'pc ' + str(frame_size) + '+' + str(inner_rad) plt.title(title, fontsize = 14) # Option to fix the y-limit @@ -350,7 +363,7 @@ def contrast_curve(cube, angle_list, psf_template, fwhm, pxscale, starphot, 'distance': rad_samp, 'noise': noise_samp_sm}) if full_output: - return (datafr, cube_fc_all, frame_fc_all, frame_nofc, fc_map_all) + return datafr, cube_fc_all, frame_fc_all, frame_nofc, fc_map_all else: return datafr @@ -359,21 +372,23 @@ def throughput(cube, angle_list, psf_template, fwhm, pxscale, algo, nbranch=1, theta=0, inner_rad=1, fc_rad_sep=3, wedge=(0,360), fc_snr=10.0, full_output=False, imlib='opencv', interpolation='lanczos4', verbose=True, **algo_dict): - """ Measures the throughput for chosen algorithm and input dataset. The - final throughput is the average of the same procedure measured in *nbranch* - azimutally equidistant branches. + """ Measures the throughput for chosen algorithm and input dataset (ADI or + ADI+IFS). The final throughput is the average of the same procedure measured + in ``nbranch`` azimutally equidistant branches. Parameters - ---------- + ---------_ cube : array_like - The input cube without fake companions. + The input cube, 2d (ADI data) or 3d array (IFS data), without fake + companions. angle_list : array_like Vector with the parallactic angles. psf_template : array_like Frame with the psf template for the fake companion(s). PSF must be centered in array. Normalization is done internally. - fwhm : float - FWHM in pixels. + fwhm: int or float or 1d array, optional + The the Full Width Half Maximum in pixels. It can handle a different + FWHM value for different wavelengths (IFS data). pxscale : float Plate scale in arcsec/px. algo : callable or function @@ -440,26 +455,52 @@ def throughput(cube, angle_list, psf_template, fwhm, pxscale, algo, nbranch=1, array = cube parangles = angle_list - if not array.ndim == 3: - raise TypeError('The input array is not a cube') - if not array.shape[0] == parangles.shape[0]: - raise TypeError('Input vector or parallactic angles has wrong length') - if not psf_template.ndim==2: - raise TypeError('Template PSF is not a frame or 2d array') - if not hasattr(algo, '__call__'): - raise TypeError('Parameter *algo* must be a callable function') - if not fc_rad_sep>=3 or not fc_rad_sep<=int((array.shape[1]/2.)/fwhm)-1: - msg = 'Too large separation between companions in the radial patterns. ' - msg += 'Should lie between 3 and {:}' - raise ValueError(msg.format(int((array.shape[1]/2.)/fwhm)-1)) - if not isinstance(inner_rad, int): - raise TypeError('inner_rad must be an integer') - angular_range = wedge[1]-wedge[0] - if nbranch>1 and angular_range<360: - msg = 'Only a single branch is allowed when working on a wedge' - raise RuntimeError(msg) - - if verbose: start_time = time_ini() + if not (array.ndim == 3 or array.ndim == 4): + raise TypeError('The input array is not a 3d or 4d cube') + else: + if array.ndim == 3: + if not array.shape[0] == parangles.shape[0]: + msg = 'Input parallactic angles vector has wrong length' + raise TypeError(msg) + if not psf_template.ndim == 2: + raise TypeError('Template PSF is not a frame or 2d array') + maxfcsep = int((array.shape[1]/2.)/fwhm)-1 + if not fc_rad_sep >= 3 or not fc_rad_sep <= maxfcsep: + msg = 'Too large separation between companions in the radial ' + msg += 'patterns. Should lie between 3 and {:}' + raise ValueError(msg.format(maxfcsep)) + + elif array.ndim == 4: + if not array.shape[1] == parangles.shape[0]: + msg = 'Input vector or parallactic angles has wrong length' + raise TypeError(msg) + if not psf_template.ndim == 3: + raise TypeError('Template PSF is not a frame, 3d array') + if 'scale_list' not in algo_dict: + raise ValueError('Vector of wavelength not found') + else: + if not algo_dict['scale_list'].shape[0] == array.shape[0]: + raise TypeError('Input wavelength vector has wrong length') + maxfcsep = int((array.shape[2] / 2.) / fwhm) - 1 + if not fc_rad_sep >= 3 or not fc_rad_sep <= maxfcsep: + msg = 'Too large separation between companions in the ' + msg += 'radial patterns. Should lie between 3 and {:}' + raise ValueError(msg.format(maxfcsep)) + + # TODO: extend to even-sized psf arrays + if psf_template.shape[1]%2 == 0: + raise ValueError("Only odd-sized PSF is accepted") + if not hasattr(algo, '__call__'): + raise TypeError('Parameter *algo* must be a callable function') + if not isinstance(inner_rad, int): + raise TypeError('inner_rad must be an integer') + angular_range = wedge[1] - wedge[0] + if nbranch > 1 and angular_range < 360: + msg = 'Only a single branch is allowed when working on a wedge' + raise RuntimeError(msg) + + if verbose: + start_time = time_ini() #*************************************************************************** # Compute noise in concentric annuli on the "empty frame" if 'cube' and 'angle_list' and 'verbose' in inspect.getargspec(algo).args: @@ -483,83 +524,202 @@ def throughput(cube, angle_list, psf_template, fwhm, pxscale, algo, nbranch=1, print('Measured annulus-wise noise in resulting frame') timing(start_time) - # We crop the PSF and check if PSF has been normalized (so that flux in - # 1*FWHM aperture = 1) and fix if needed - psf_template = psf_norm(psf_template, size=3*fwhm, fwhm=fwhm) - - #*************************************************************************** - # Initialize the fake companions - angle_branch = angular_range / nbranch - # signal-to-noise ratio of injected fake companions - snr_level = fc_snr * np.ones_like(noise) - - thruput_arr = np.zeros((nbranch, noise.shape[0])) - fc_map_all = np.zeros((nbranch*fc_rad_sep, array.shape[1], array.shape[2])) - frame_fc_all = fc_map_all.copy() - cube_fc_all = np.zeros((nbranch*fc_rad_sep, array.shape[0], array.shape[1], - array.shape[2])) - cy, cx = frame_center(array[0]) - - # each branch is computed separately - for br in range(nbranch): - # each pattern is computed separately. For each pattern the companions - # are separated by "fc_rad_sep * fwhm", interleaving the injections - for irad in range(fc_rad_sep): - radvec = vector_radd[irad::fc_rad_sep] - cube_fc = array.copy() - # filling map with small numbers - fc_map = np.ones_like(array[0]) * 1e-6 - fcy = []; fcx = [] - for i in range(radvec.shape[0]): - flux = snr_level[irad+i*fc_rad_sep] * noise[irad+i*fc_rad_sep] - cube_fc = cube_inject_companions(cube_fc, psf_template, - parangles, flux, pxscale, - rad_dists=[radvec[i]], - theta=br*angle_branch + theta, - imlib=imlib, verbose=False, - interpolation=interpolation) - y = cy + radvec[i] * np.sin(np.deg2rad(br*angle_branch + theta)) - x = cx + radvec[i] * np.cos(np.deg2rad(br*angle_branch + theta)) - fc_map = frame_inject_companion(fc_map, psf_template, y, x, - flux, imlib, interpolation) - fcy.append(y); fcx.append(x) - - if verbose: - msg2 = 'Fake companions injected in branch {:} (pattern {:}/{:})' - print(msg2.format(br+1, irad+1, fc_rad_sep)) - timing(start_time) - - #******************************************************************* - if 'cube' and 'angle_list' and 'verbose' in inspect.getargspec(algo).args: - if 'fwhm' in inspect.getargspec(algo).args: - frame_fc = algo(cube=cube_fc, angle_list=parangles, - fwhm=fwhm, verbose=False, **algo_dict) + if cube.ndim == 3: + # We crop the PSF and check if PSF has been normalized (so that flux in + # 1*FWHM aperture = 1) and fix if needed + # TODO: verify it handles odd and even-sized cases + new_psf_size = 3*fwhm + if new_psf_size%2 == 0: + new_psf_size += 1 + psf_template = psf_norm(psf_template, size=min(new_psf_size, + psf_template.shape[1]), + fwhm=fwhm) + + #*********************************************************************** + # Initialize the fake companions + angle_branch = angular_range / nbranch + # signal-to-noise ratio of injected fake companions + snr_level = fc_snr * np.ones_like(noise) + + thruput_arr = np.zeros((nbranch, noise.shape[0])) + fc_map_all = np.zeros((nbranch * fc_rad_sep, array.shape[1], + array.shape[2])) + frame_fc_all = fc_map_all.copy() + cube_fc_all = np.zeros((nbranch * fc_rad_sep, array.shape[0], + array.shape[1], array.shape[2])) + cy, cx = frame_center(array[0]) + + # each branch is computed separately + for br in range(nbranch): + # each pattern is computed separately. For each one the companions + # are separated by "fc_rad_sep * fwhm", interleaving the injections + for irad in range(fc_rad_sep): + radvec = vector_radd[irad::fc_rad_sep] + cube_fc = array.copy() + # filling map with small numbers + fc_map = np.ones_like(array[0]) * 1e-6 + fcy = [] + fcx = [] + for i in range(radvec.shape[0]): + flux = snr_level[irad+i*fc_rad_sep] * noise[irad + i * + fc_rad_sep] + cube_fc = cube_inject_companions(cube_fc, psf_template, + parangles, flux, pxscale, + rad_dists=[radvec[i]], + theta=br*angle_branch + + theta, + imlib=imlib, verbose=False, + interpolation=interpolation) + y = cy + radvec[i] * np.sin(np.deg2rad(br * angle_branch + + theta)) + x = cx + radvec[i] * np.cos(np.deg2rad(br * angle_branch + + theta)) + fc_map = frame_inject_companion(fc_map, psf_template, y, x, + flux, imlib, interpolation) + fcy.append(y) + fcx.append(x) + + if verbose: + msg2 = 'Fake companions injected in branch {:} ' + msg2 += '(pattern {:}/{:})' + print(msg2.format(br+1, irad+1, fc_rad_sep)) + timing(start_time) + + #*************************************************************** + if 'cube' and 'angle_list' and 'verbose' in inspect.getargspec(algo).args: + if 'fwhm' in inspect.getargspec(algo).args: + frame_fc = algo(cube=cube_fc, angle_list=parangles, + fwhm=fwhm, verbose=False, **algo_dict) + else: + frame_fc = algo(cube=cube_fc, angle_list=parangles, + verbose=False, **algo_dict) else: - frame_fc = algo(cube=cube_fc, angle_list=parangles, - verbose=False, **algo_dict) - - if verbose: - msg3 = 'Cube with fake companions processed with {:}' - msg3 += '\nMeasuring its annulus-wise throughput' - print(msg3.format(algo.__name__)) - timing(start_time) - - #******************************************************************* - injected_flux = aperture_flux(fc_map, fcy, fcx, fwhm, ap_factor=1, - mean=False, verbose=False) - recovered_flux = aperture_flux((frame_fc - frame_nofc), fcy, fcx, - fwhm, ap_factor=1, mean=False, - verbose=False) - thruput = (recovered_flux)/injected_flux - thruput[np.where(thruput<0)] = 0 - - thruput_arr[br, irad::fc_rad_sep] = thruput - fc_map_all[br*fc_rad_sep+irad, :, :] = fc_map - frame_fc_all[br*fc_rad_sep+irad, :, :] = frame_fc - cube_fc_all[br*fc_rad_sep+irad, :, :, :] = cube_fc + msg = 'Input algorithm must have at least 3 parameters: ' + msg += 'cube, angle_list and versbose' + raise ValueError(msg) + + if verbose: + msg3 = 'Cube with fake companions processed with {:}' + msg3 += '\nMeasuring its annulus-wise throughput' + print(msg3.format(algo.__name__)) + timing(start_time) + + #*************************************************************** + injected_flux = aperture_flux(fc_map, fcy, fcx, fwhm, + ap_factor=1, mean=False, + verbose=False) + recovered_flux = aperture_flux((frame_fc - frame_nofc), fcy, + fcx, fwhm, ap_factor=1, + mean=False, verbose=False) + thruput = recovered_flux / injected_flux + thruput[np.where(thruput < 0)] = 0 + + thruput_arr[br, irad::fc_rad_sep] = thruput + fc_map_all[br*fc_rad_sep+irad, :, :] = fc_map + frame_fc_all[br*fc_rad_sep+irad, :, :] = frame_fc + cube_fc_all[br*fc_rad_sep+irad, :, :, :] = cube_fc + + elif cube.ndim == 4: + if isinstance(fwhm, (int, float)): + fwhm = [fwhm for _ in range(array.shape[0])] + + # We crop the PSF and check if PSF has been normalized (so that flux in + # 1*FWHM aperture = 1) and fix if needed + new_psf_size = 3 * np.max(fwhm) + if new_psf_size%2 == 0: + new_psf_size += 1 + + psf_template = psf_norm(psf_template, fwhm=fwhm, + size=min(new_psf_size, psf_template.shape[1])) + + # ********************************************************************** + # Initialize the fake companions + angle_branch = angular_range / nbranch + # signal-to-noise ratio of injected fake companions + snr_level = 10.0 * np.ones_like(noise) + + thruput_arr = np.zeros((nbranch, noise.shape[0])) + fc_map_all = np.zeros((nbranch * fc_rad_sep, array.shape[0], + array.shape[2], array.shape[3])) + frame_fc_all = fc_map_all.copy() + cube_fc_all = np.zeros( + (nbranch * fc_rad_sep, array.shape[0], array.shape[1], + array.shape[2], array.shape[3])) + cy, cx = frame_center(array[0, 0]) + + # each branch is computed separately + for br in range(nbranch): + # each pattern is computed separately. For each pattern the + # companions are separated by "fc_rad_sep * fwhm" + # radius = vector_radd[irad::fc_rad_sep] + for irad in range(fc_rad_sep): + radvec = vector_radd[irad::fc_rad_sep] + thetavec = range(int(theta), int(theta) + 360, + int(360 / len(radvec))) + cube_fc = array.copy() + # filling map with small numbers + fc_map = np.ones_like(array[:, 0]) * 1e-6 + fcy = [] + fcx = [] + for i in range(radvec.shape[0]): + flux = snr_level[irad + i * fc_rad_sep] * noise[irad + i * \ + fc_rad_sep] + cube_fc = cube_inject_companions(cube_fc, psf_template, + parangles, flux, pxscale, + rad_dists=[radvec[i]], + theta=thetavec[i], + verbose=False) + y = cy + radvec[i] * np.sin(np.deg2rad(br * angle_branch + + thetavec[i])) + x = cx + radvec[i] * np.cos(np.deg2rad(br * angle_branch + + thetavec[i])) + fc_map = frame_inject_companion(fc_map, psf_template, y, x, + flux) + fcy.append(y) + fcx.append(x) + + if verbose: + msg2 = 'Fake companions injected in branch {:} ' + msg2 += '(pattern {:}/{:})' + print(msg2.format(br + 1, irad + 1, fc_rad_sep)) + timing(start_time) + + # ************************************************************** + if 'cube' and 'angle_list' and 'verbose' in inspect.getargspec( + algo).args: + if 'fwhm' in inspect.getargspec(algo).args: + frame_fc = algo(cube=cube_fc, angle_list=parangles, + fwhm=fwhm, verbose=False, **algo_dict) + else: + frame_fc = algo(cube=cube_fc, angle_list=parangles, + verbose=False, **algo_dict) + + if verbose: + msg3 = 'Cube with fake companions processed with {:}' + msg3 += '\nMeasuring its annulus-wise throughput' + print(msg3.format(algo.__name__)) + timing(start_time) + + # ************************************************************** + injected_flux = [aperture_flux(fc_map[i], fcy, fcx, fwhm[i], + ap_factor=1, mean=False, + verbose=False) for i in + range(array.shape[0])] + injected_flux = np.mean(injected_flux, axis=0) + recovered_flux = aperture_flux((frame_fc - frame_nofc), fcy, + fcx, np.mean(fwhm), ap_factor=1, + mean=False, verbose=False) + thruput = recovered_flux / injected_flux + thruput[np.where(thruput < 0)] = 0 + + thruput_arr[br, irad::fc_rad_sep] = thruput + fc_map_all[br * fc_rad_sep + irad, :, :] = fc_map + frame_fc_all[br * fc_rad_sep + irad, :, :] = frame_fc + cube_fc_all[br * fc_rad_sep + irad, :, :, :, :] = cube_fc if verbose: - print('Finished measuring the throughput in {:} branches'.format(nbranch)) + msg = 'Finished measuring the throughput in {:} branches' + print(msg.format(nbranch)) timing(start_time) if full_output: @@ -569,7 +729,6 @@ def throughput(cube, angle_list, psf_template, fwhm, pxscale, algo, nbranch=1, return thruput_arr, vector_radd - def noise_per_annulus(array, separation, fwhm, init_rad=None, wedge=(0,360), verbose=False, debug=False): """ Measures the noise as the standard deviation of apertures defined in diff --git a/vip_hci/phot/fakecomp.py b/vip_hci/phot/fakecomp.py index 83585236..9e0024e3 100644 --- a/vip_hci/phot/fakecomp.py +++ b/vip_hci/phot/fakecomp.py @@ -21,12 +21,15 @@ from ..var import frame_center, fit_2dgaussian, get_circle -### TODO: remove this in VIP v1.0.0. Created for backward compatibility +# TODO: remove this in VIP v1.0.0. Created for backward compatibility def inject_fcs_cube(array, psf_template, angle_list, flevel, plsc, rad_dists, n_branches=1, theta=0, imlib='opencv', verbose=True): return cube_inject_companions(array, psf_template, angle_list, flevel, plsc, rad_dists, n_branches, theta, imlib, verbose) + + +# TODO: remove this in VIP v1.0.0. Created for backward compatibility def inject_fc_frame(array, array_fc, pos_y, pos_x, flux): return frame_inject_companion(array, array_fc, pos_y, pos_x, flux) @@ -70,20 +73,22 @@ def cube_inject_companions(array, psf_template, angle_list, flevel, plsc, """ - if not (array.ndim==3 or array.ndim==4): + if not (array.ndim == 3 or array.ndim == 4): raise ValueError('Array is not a cube, 3d or 4d array') - if array.ndim==4: - if not psf_template.ndim==3: + if array.ndim == 4: + if not psf_template.ndim == 3: raise ValueError('Psf must be a 3d array') # ADI case - if array.ndim==3: - + if array.ndim == 3: + # TODO: extend to even-sized psf arrays + if psf_template.shape[1]%2 == 0: + raise ValueError("Only odd-sized PSF is accepted") ceny, cenx = frame_center(array[0]) ceny = int(float(ceny)) cenx = int(float(cenx)) rad_dists = np.array(rad_dists) - if not rad_dists[-1]1.1 or fwhm_flux<0.9: - psf_norm_array = psfs/np.array(fwhm_aper_phot['aperture_sum']) - else: - psf_norm_array = psfs - - if threshold is not None: - psf_norm_array[np.where(psf_norm_array < threshold)] = 0 - - if mask_core is not None: - psf_norm_array = get_circle(psf_norm_array, radius=mask_core) - - if full_output: - return psf_norm_array, fwhm_flux - else: - return psf_norm_array + psfs = frame_shift(psfs, -shifty, -shiftx, imlib=imlib, + interpolation=interpolation) + for _ in range(2): + centroidy, centroidx = fit_2dgaussian(psfs, fwhmx=fwhm, + fwhmy=fwhm) + cy, cx = frame_center(psfs, verbose=False) + shiftx, shifty = centroidx - cx, centroidy - cy + psfs = frame_shift(psfs, -shifty, -shiftx) + + # we check whether the flux is normalized and fix it if needed + fwhm_aper = photutils.CircularAperture((frame_center(psfs)), fwhm/2.) + fwhm_aper_phot = photutils.aperture_photometry(psfs, fwhm_aper, + method='exact') + fwhm_flux = np.array(fwhm_aper_phot['aperture_sum']) + if verbose: + print("Flux in 1xFWHM aperture: {}".format(fwhm_flux)) + + if fwhm_flux > 1.1 or fwhm_flux < 0.9: + psf_norm_array = psfs / np.array(fwhm_aper_phot['aperture_sum']) + else: + psf_norm_array = psfs + + if threshold is not None: + psf_norm_array[np.where(psf_norm_array < threshold)] = 0 + + if mask_core is not None: + psf_norm_array = get_circle(psf_norm_array, radius=mask_core) + + if full_output: + return psf_norm_array, fwhm_flux + else: + return psf_norm_array + #--------------------------------------------------------------------------- + + if array.ndim == 2: + res = psf_norm_2d(array, fwhm, size, threshold, mask_core, imlib, + interpolation, full_output, verbose) + + elif array.ndim == 3: + if isinstance(fwhm, (int, float)): + fwhm = [fwhm for _ in range(array.shape[0])] + array_out = np.zeros((array.shape[0], size, size)) + if full_output: + fwhm_flux = np.zeros((array.shape[0])) + + for fr in range(array.shape[0]): + restemp = psf_norm_2d(array[fr], fwhm[fr], size, threshold, + mask_core, imlib, interpolation, full_output, + False) + if full_output: + array_out[fr] = restemp[0] + fwhm_flux[fr] = restemp[1] + else: + array_out[fr] = restemp + if full_output: + res = (array_out, fwhm_flux) + else: + res = array_out - return psf_norm + return res def _cube_shift(cube, y, x, imlib, interpolation): diff --git a/vip_hci/phot/snr.py b/vip_hci/phot/snr.py index 001a60f1..e8eb5034 100644 --- a/vip_hci/phot/snr.py +++ b/vip_hci/phot/snr.py @@ -1,7 +1,7 @@ #! /usr/bin/env python """ -Module with SNR calculation functions. +Module with S/N calculation functions. """ from __future__ import division @@ -83,14 +83,14 @@ def snrmap(array, fwhm, plot=False, mode='sss', source_mask=None, nproc=None, if source_mask is None: pool = Pool(processes=int(nproc)) - res = pool.map(eval_func_tuple, itt.izip(itt.repeat(F),itt.repeat(array), - coords, itt.repeat(fwhm), - itt.repeat(True))) + res = pool.map(eval_func_tuple, zip(itt.repeat(F),itt.repeat(array), + coords, itt.repeat(fwhm), + itt.repeat(True))) res = np.array(res) pool.close() - yy = res[:,0] - xx = res[:,1] - snr = res[:,2] + yy = res[:, 0] + xx = res[:, 1] + snr = res[:, 2] snrmap[yy.astype('int'), xx.astype('int')] = snr else: # checking the mask with the sources @@ -131,9 +131,9 @@ def snrmap(array, fwhm, plot=False, mode='sss', source_mask=None, nproc=None, coor_rest = [(y,x) for (y,x) in zip(yy, xx) if (y,x) not in coor_ann] pool1 = Pool(processes=int(nproc)) - res = pool1.map(eval_func_tuple, itt.izip(itt.repeat(F),itt.repeat(array), - coor_rest, itt.repeat(fwhm), - itt.repeat(True))) + res = pool1.map(eval_func_tuple, zip(itt.repeat(F),itt.repeat(array), + coor_rest, itt.repeat(fwhm), + itt.repeat(True))) res = np.array(res) pool1.close() yy = res[:,0] @@ -142,10 +142,10 @@ def snrmap(array, fwhm, plot=False, mode='sss', source_mask=None, nproc=None, snrmap[yy.astype('int'), xx.astype('int')] = snr pool2 = Pool(processes=int(nproc)) - res = pool2.map(eval_func_tuple, itt.izip(itt.repeat(F), - itt.repeat(array_sources), - coor_ann, itt.repeat(fwhm), - itt.repeat(True))) + res = pool2.map(eval_func_tuple, zip(itt.repeat(F), + itt.repeat(array_sources), + coor_ann, itt.repeat(fwhm), + itt.repeat(True))) res = np.array(res) pool2.close() yy = res[:,0] @@ -191,8 +191,9 @@ def snrmap_fast(array, fwhm, nproc=None, plot=False, verbose=True): Frame with the same size as the input frame with each pixel. """ - if verbose: start_time = time_ini() - if not array.ndim==2: + if verbose: + start_time = time_ini() + if not array.ndim == 2: raise TypeError('Input array is not a 2d array or image.') cy, cx = frame_center(array) @@ -205,41 +206,42 @@ def snrmap_fast(array, fwhm, nproc=None, plot=False, verbose=True): mask = get_annulus(array, (fwhm/2)+1, width-1) mask = np.ma.make_mask(mask) yy, xx = np.where(mask) - coords = [(x,y) for (x,y) in zip(xx,yy)] + coords = [(x, y) for (x, y) in zip(xx, yy)] if nproc is None: nproc = int((cpu_count()/2)) # Hyper-threading doubles the # of cores - if nproc==1: - for y,x in zip(yy,xx): - snrmap[y,x] = _snr_approx(array, (x,y), fwhm, cy, cx)[2] - elif nproc>1: + if nproc == 1: + for y,x in zip(yy, xx): + snrmap[y,x] = _snr_approx(array, (x, y), fwhm, cy, cx)[2] + elif nproc > 1: pool = Pool(processes=int(nproc)) - res = pool.map(eval_func_tuple, itt.izip(itt.repeat(_snr_approx), - itt.repeat(array), - coords, itt.repeat(fwhm), - itt.repeat(cy), - itt.repeat(cx))) + res = pool.map(eval_func_tuple, zip(itt.repeat(_snr_approx), + itt.repeat(array), coords, + itt.repeat(fwhm), itt.repeat(cy), + itt.repeat(cx))) res = np.array(res) pool.close() - yy = res[:,0] - xx = res[:,1] - snr = res[:,2] + yy = res[:, 0] + xx = res[:, 1] + snr = res[:, 2] snrmap[yy.astype('int'), xx.astype('int')] = snr - if plot: pp_subplots(snrmap, colorb=True, title='SNRmap') + if plot: + pp_subplots(snrmap, colorb=True, title='SNRmap') if verbose: print("S/N map created using {:} processes.".format(nproc)) timing(start_time) return snrmap -def _snr_approx(array, xxx_todo_changeme, fwhm, centery, centerx): + +def _snr_approx(array, source_xy, fwhm, centery, centerx): """ array - frame convolved with top hat kernel """ - (sourcex,sourcey) = xxx_todo_changeme - rad = dist(centery,centerx,sourcey,sourcex) + sourcex, sourcey = source_xy + rad = dist(centery, centerx, sourcey, sourcex) ind_aper = draw.circle(sourcey, sourcex, fwhm/2.) # noise : STDDEV in convolved array of 1px wide annulus (while # masking the flux aperture) * correction of # of resolution elements @@ -253,10 +255,9 @@ def _snr_approx(array, xxx_todo_changeme, fwhm, centery, centerx): snr = signal / noise return sourcey, sourcex, snr - - -def snr_ss(array, xxx_todo_changeme1, fwhm, out_coor=False, plot=False, - verbose=False, full_output=False): + +def snr_ss(array, source_xy, fwhm, out_coor=False, plot=False, verbose=False, + full_output=False): # Leave the order of parameters as it is, the same for both snr functions # to be compatible with the snrmap parallel implementation """Calculates the SNR (signal to noise ratio) of a single planet in a @@ -291,8 +292,7 @@ def snr_ss(array, xxx_todo_changeme1, fwhm, out_coor=False, plot=False, sourcey, sourcex, f_source, fluxes.std(), snr """ - (source_xy) = xxx_todo_changeme1 - if not array.ndim==2: + if not array.ndim == 2: raise TypeError('Input array is not a frame or 2d array') if out_coor and full_output: raise TypeError('One of the 2 must be False') @@ -313,10 +313,10 @@ def snr_ss(array, xxx_todo_changeme1, fwhm, out_coor=False, plot=False, xx[0] = sourcex - centerx yy[0] = sourcey - centery for i in range(number_apertures-1): - if sens=='clock': + if sens == 'clock': xx[i+1] = cosangle*xx[i] + sinangle*yy[i] yy[i+1] = cosangle*yy[i] - sinangle*xx[i] - elif sens=='counterclock': + elif sens == 'counterclock': xx[i+1] = cosangle*xx[i] - sinangle*yy[i] yy[i+1] = cosangle*yy[i] + sinangle*xx[i] @@ -367,7 +367,7 @@ def snr_ss(array, xxx_todo_changeme1, fwhm, out_coor=False, plot=False, return snr -def snr_peakstddev(array, xxx_todo_changeme2, fwhm, out_coor=False, plot=False, +def snr_peakstddev(array, source_xy, fwhm, out_coor=False, plot=False, verbose=False): """Calculates the S/N (signal to noise ratio) of a single planet in a post-processed (e.g. by LOCI or PCA) frame. The signal is taken as the ratio @@ -397,8 +397,7 @@ def snr_peakstddev(array, xxx_todo_changeme2, fwhm, out_coor=False, plot=False, Value of the SNR for the given planet or test speckle. """ - (source_xy) = xxx_todo_changeme2 - sourcex, sourcey = source_xy + sourcex, sourcey = source_xy centery, centerx = frame_center(array) rad = dist(centery,centerx,sourcey,sourcex) @@ -425,7 +424,7 @@ def snr_peakstddev(array, xxx_todo_changeme2, fwhm, out_coor=False, plot=False, fill=False) ax.add_patch(circ2) aper = plt.Circle((sourcex, sourcey), radius=fwhm/2., color='b', - fill=False) # Coordinates (X,Y) + fill=False) # Coordinates (X,Y) ax.add_patch(aper) plt.show() plt.gray() @@ -433,8 +432,4 @@ def snr_peakstddev(array, xxx_todo_changeme2, fwhm, out_coor=False, plot=False, if out_coor: return sourcey, sourcex, snr else: - return snr - - - - + return snr \ No newline at end of file diff --git a/vip_hci/preproc/__init__.py b/vip_hci/preproc/__init__.py index 4ad1a55f..41cfa040 100644 --- a/vip_hci/preproc/__init__.py +++ b/vip_hci/preproc/__init__.py @@ -30,7 +30,6 @@ from .badframes import * from .badpixremoval import * from .cosmetics import * -from .cosmetics_ifs import * from .derotation import * from .parangles import * from .recentering import * diff --git a/vip_hci/preproc/cosmetics.py b/vip_hci/preproc/cosmetics.py index b42ea232..0f3a958e 100644 --- a/vip_hci/preproc/cosmetics.py +++ b/vip_hci/preproc/cosmetics.py @@ -8,32 +8,35 @@ from __future__ import division from __future__ import print_function -__author__ = 'Carlos Alberto Gomez Gonzalez' +__author__ = 'Carlos Alberto Gomez Gonzalez, V. Christiaens @ UChile/ULg' __all__ = ['cube_crop_frames', 'cube_drop_frames', - 'frame_crop'] + 'frame_crop', + 'cube_correct_nan', + 'approx_stellar_position'] import numpy as np +from astropy.stats import sigma_clipped_stats +from ..stats import sigma_filter from ..var import frame_center, get_square -def cube_crop_frames(array, size, xy=None, verbose=True, full_output = False): - """Crops frames in a cube (3d array). If size is an even value it'll be - increased by one to make it odd. +def cube_crop_frames(array, size, xy=None, verbose=True, full_output=False): + """Crops frames in a cube (3d or 4d array). Parameters ---------- array : array_like - Input 3d array, cube. + Input 3d or 4d array. size : int Size of the desired central sub-array in each frame. xy : tuple of ints - X, Y coordinates of new frame center. If you are getting the + X, Y coordinates of new frame center. If you are getting the coordinates from ds9 subtract 1, python has 0-based indexing. - verbose : {True, False}, bool optional + verbose : bool optional If True message of completion is showed. - full_output: {False, True} + full_output: bool optional If true, returns cenx and ceny in addition to array_view Returns @@ -42,38 +45,60 @@ def cube_crop_frames(array, size, xy=None, verbose=True, full_output = False): Cube with cropped frames. """ - if not array.ndim == 3: - raise TypeError('Array is not 3d, not a cube') + if not (array.ndim == 3 or array.ndim == 4): + raise TypeError('Array is not a cube, 3d or 4d array') if not isinstance(size, int): raise TypeError('Size must be integer') - - if size%2!=0: size -= 1 - if size >= array.shape[1]: - raise ValueError('The crop size is equal or bigger than the frame size') - - # wing is added to the sides of the subframe center. - wing = int(size/2) - + if array.shape[2]%2 == 0: # assuming square frames, both 3d or 4d case + if size%2 != 0: + msg = "Warning: the new size is odd and the original frame size is" + msg += " even. Make sure you are setting properly xy parameter" + print(msg) + else: + if size % 2 == 0: + msg = "Warning: the new size is even and the original frame size is" + msg += " odd. Make sure you are setting properly xy parameter" + print(msg) if xy is not None: if not (isinstance(xy[0], int) or isinstance(xy[1], int)): raise TypeError('XY must be a tuple of integers') - cenx, ceny = xy - # Note the +1 when closing the interval (python doesn't include the - # endpoint) - array_view = array[:,int(ceny-wing):int(ceny+wing+1), - int(cenx-wing):int(cenx+wing+1)].copy() + if size >= array.shape[2]: # assuming square frames, both 3d or 4d case + raise ValueError('The new size is equal or bigger than the frame size') + + # wing is added to the sides of the subframe center + if size%2 != 0: + wing = int(np.floor(size / 2.)) + else: + wing = (size / 2.) - 0.5 + + if array.ndim == 3: + if xy is not None: + cenx, ceny = xy + else: + ceny, cenx = frame_center(array[0], verbose=False) + + # +1 because python doesn't include the endpoint when slicing + array_view = array[:, int(ceny-wing):int(ceny+wing+1), + int(cenx-wing):int(cenx+wing+1)] if verbose: - print() - msg = "Cube cropped; new size [{:},{:},{:}] centered at ({:},{:})" - print(msg.format(array.shape[0], size+1, size+1, cenx, ceny)) - else: - cy, cx = frame_center(array[0], verbose=False) - array_view = array[:, int(cy-wing):int(cy+wing+1), - int(cx-wing):int(cx+wing+1)].copy() + msg = "\nCube cropped; new size [{:},{:},{:}] centered at ({:},{:})" + print(msg.format(array_view.shape[0], array_view.shape[1], + array_view.shape[2], cenx, ceny)) + + elif array.ndim == 4: + if xy is not None: + cenx, ceny = xy + else: + ceny, cenx = frame_center(array[0, 0], verbose=False) + + # +1 because python doesn't include the endpoint when slicing + array_view = array[:, :, int(ceny - wing):int(ceny + wing + 1), + int(cenx - wing):int(cenx + wing + 1)] if verbose: - print() - msg = "Cube cropped with new size [{:},{:},{:}]" - print(msg.format(array.shape[0], size+1, size+1)) + msg = "\nCube cropped; new size [{:},{:},{:},{:}] centered at ({:},{:})" + print(msg.format(array_view.shape[0], array_view.shape[1], + array_view.shape[2], array_view.shape[3], + cenx, ceny)) # Option to return the cenx and ceny for double checking that the frame crop # did not exceed the limit @@ -84,8 +109,7 @@ def cube_crop_frames(array, size, xy=None, verbose=True, full_output = False): def frame_crop(array, size, cenxy=None, verbose=True): - """Crops a square frame (2d array). Wrapper of function *vortex.var.get_square*. - If size is an even value it'll be increased by one to make it odd. + """Crops a square frame (2d array). Parameters ---------- @@ -95,7 +119,7 @@ def frame_crop(array, size, cenxy=None, verbose=True): Size of the subframe. cenxy : tuple, optional Coordinates of the center of the subframe. - verbose : {True, False}, bool optional + verbose : bool optional If True message of completion is showed. Returns @@ -104,10 +128,11 @@ def frame_crop(array, size, cenxy=None, verbose=True): Sub array. """ - if not array.ndim==2: + if not array.ndim == 2: raise TypeError('Array is not a frame or 2d array') - if size>=array.shape[0]: - raise RuntimeError('Cropping size if equal or larger than the original size') + if size >= array.shape[0]: + msg = 'Cropping size is equal or larger than the original size' + raise RuntimeError(msg) if not cenxy: ceny, cenx = frame_center(array, verbose=False) @@ -116,36 +141,34 @@ def frame_crop(array, size, cenxy=None, verbose=True): array_view = get_square(array, size, ceny, cenx) if verbose: - print() - print("Done frame cropping") + print("\nDone frame cropping") return array_view def cube_drop_frames(array, n, m): """Discards frames at the beginning or the end of a cube (axis 0). - + Parameters ---------- - array : array_like + array : array_like Input 3d array, cube. n : int Index of the first frame to be kept. Frames before this one are dropped. m : int Index of the last frame to be kept. Frames after this one are dropped. - + Returns ------- array_view : array_like Cube with new size (axis 0). - + """ if m>array.shape[0]: raise TypeError('End index must be smaller than the # of frames') - + array_view = array[n:m+1, :, :].copy() - print() print("Done discarding frames from cube") - return array_view + return array_view def frame_remove_stripes(array): @@ -159,3 +182,188 @@ def frame_remove_stripes(array): array[:,i] = array[:,i] - mean[i] return array + +def cube_correct_nan(cube, neighbor_box=3, min_neighbors=3, verbose=False, + half_res_y=False): + """Sigma filtering of nan pixels in a whole frame or cube. Intended for + SINFONI data. + + Parameters + ---------- + cube : cube_like + Input 3d or 2d array. + neighbor_box : int, optional + The side of the square window around each pixel where the sigma and + median are calculated for the nan pixel correction. + min_neighbors : int, optional + Minimum number of good neighboring pixels to be able to correct the + bad/nan pixels. + verbose: {False,True} bool, optional + Whether to print more information or not during processing + half_res_y: bool, {True,False}, optional + Whether the input data have every couple of 2 rows identical, i.e. there + is twice less angular resolution vertically than horizontally (e.g. + SINFONI data). The algorithm goes twice faster if this option is + rightfully set to True. + + Returns + ------- + obj_tmp : array_like + Output cube with corrected nan pixels in each frame + """ + + obj_tmp = cube.copy() + + ndims = obj_tmp.ndim + if ndims != 2 and ndims != 3: + raise TypeError("Input object is not two or three dimensional") + + if neighbor_box < 3 or neighbor_box % 2 == 0: + raise ValueError('neighbor_box should be an odd value greater than 3') + max_neigh = sum(range(3, neighbor_box + 2, 2)) + if min_neighbors > max_neigh: + min_neighbors = max_neigh + msg = "Warning! min_neighbors was reduced to " + str(max_neigh) + \ + " to avoid bugs. \n" + print(msg) + + def nan_corr_2d(obj_tmp): + n_x = obj_tmp.shape[1] + n_y = obj_tmp.shape[0] + + if half_res_y: + if n_y % 2 != 0: + msg = 'The input frames do not have an even number of rows. ' + msg2 = 'Hence, you should probably not be using the option ' + msg3 = 'half_res_y = True.' + raise ValueError(msg + msg2 + msg3) + n_y = int(n_y / 2) + frame = obj_tmp + obj_tmp = np.zeros([n_y, n_x]) + for yy in range(n_y): + obj_tmp[yy] = frame[2 * yy] + + # tuple with the 2D indices of each nan value of the frame + nan_indices = np.where(np.isnan(obj_tmp)) + nan_map = np.zeros_like(obj_tmp) + nan_map[nan_indices] = 1 + nnanpix = int(np.sum(nan_map)) + # Correct nan with iterative sigma filter + obj_tmp = sigma_filter(obj_tmp, nan_map, neighbor_box=neighbor_box, + min_neighbors=min_neighbors, verbose=verbose) + if half_res_y: + frame = obj_tmp + n_y = 2 * n_y + obj_tmp = np.zeros([n_y, n_x]) + for yy in range(n_y): + obj_tmp[yy] = frame[int(yy / 2)] + + return obj_tmp, nnanpix + + if ndims == 2: + obj_tmp, nnanpix = nan_corr_2d(obj_tmp) + if verbose: + print('\n There were ', nnanpix, ' nan pixels corrected.') + + elif ndims == 3: + n_z = obj_tmp.shape[0] + for zz in range(n_z): + obj_tmp[zz], nnanpix = nan_corr_2d(obj_tmp[zz]) + if verbose: + msg = 'In channel ' + str(zz) + ', there were ' + str(nnanpix) + msg2 = ' nan pixels corrected.' + print(msg + msg2) + + if verbose: + print('All nan pixels are corrected.') + + return obj_tmp + + +def approx_stellar_position(cube, fwhm, return_test=False, verbose=False): + """FIND THE APPROX COORDS OF THE STAR IN EACH CHANNEL (even the ones + dominated by noise) + + Parameters + ---------- + obj_tmp : array_like + Input 3d cube + fwhm : float or array 1D + Input full width half maximum value of the PSF for each channel. + This will be used as the standard deviation for Gaussian kernel + of the Gaussian filtering. + If float, it is assumed the same for all channels. + return_test: bool, {False,True}, optional + Whether the test result vector (a bool vector with whether the star + centroid could be find in the corresponding channel) should be returned + as well, along with the approx stellar coordinates. + verbose: {True, False}, bool optional + Chooses whether to print some additional information. + + Returns: + -------- + Array of y and x approx coordinates of the star in each channel of the cube + if return_test: it also returns the test result vector + """ + from ..phot import peak_coordinates + + obj_tmp = cube.copy() + n_z = obj_tmp.shape[0] + + if isinstance(fwhm, float) or isinstance(fwhm, int): + fwhm_scal = fwhm + fwhm = np.zeros((n_z)) + fwhm[:] = fwhm_scal + + # 1/ Write a 2-columns array with indices of all max pixel values in the cube + star_tmp_idx = np.zeros([n_z, 2]) + star_approx_idx = np.zeros([n_z, 2]) + test_result = np.ones(n_z) + for zz in range(n_z): + star_tmp_idx[zz] = peak_coordinates(obj_tmp[zz], fwhm[zz]) + + # 2/ Detect the outliers in each column + _, med_y, stddev_y = sigma_clipped_stats(star_tmp_idx[:, 0], sigma=2.5) + _, med_x, stddev_x = sigma_clipped_stats(star_tmp_idx[:, 1], sigma=2.5) + lim_inf_y = med_y - 3 * stddev_y + lim_sup_y = med_y + 3 * stddev_y + lim_inf_x = med_x - 3 * stddev_x + lim_sup_x = med_x + 3 * stddev_x + + if verbose: + print("median y of star - 3sigma = ", lim_inf_y) + print("median y of star + 3sigma = ", lim_sup_y) + print("median x of star - 3sigma = ", lim_inf_x) + print("median x of star + 3sigma = ", lim_sup_x) + + for zz in range(n_z): + if ((star_tmp_idx[zz, 0] < lim_inf_y) or ( + star_tmp_idx[zz, 0] > lim_sup_y) or + (star_tmp_idx[zz, 1] < lim_inf_x) or ( + star_tmp_idx[zz, 1] > lim_sup_x)): + test_result[zz] = 0 + + # 3/ Replace by the median of neighbouring good coordinates if need be + for zz in range(n_z): + if test_result[zz] == 0: + ii = 1 + inf_neigh = max(0, zz - ii) + sup_neigh = min(n_z - 1, zz + ii) + while test_result[inf_neigh] == 0 and test_result[sup_neigh] == 0: + ii = ii + 1 + inf_neigh = max(0, zz - ii) + sup_neigh = min(n_z - 1, zz + ii) + if test_result[inf_neigh] == 1 and test_result[sup_neigh] == 1: + star_approx_idx[zz] = np.floor((star_tmp_idx[sup_neigh] + \ + star_tmp_idx[inf_neigh]) / 2.) + elif test_result[inf_neigh] == 1: + star_approx_idx[zz] = star_tmp_idx[inf_neigh] + else: + star_approx_idx[zz] = star_tmp_idx[sup_neigh] + else: + star_approx_idx[zz] = star_tmp_idx[zz] + + if return_test: + return star_approx_idx, test_result.astype(bool) + else: + return star_approx_idx diff --git a/vip_hci/preproc/cosmetics_ifs.py b/vip_hci/preproc/cosmetics_ifs.py deleted file mode 100644 index 530a73db..00000000 --- a/vip_hci/preproc/cosmetics_ifs.py +++ /dev/null @@ -1,201 +0,0 @@ -#! /usr/bin/env python - -""" -Module with cube cosmetic functions for SDI datasets. -""" - -from __future__ import division -from __future__ import print_function - -__author__ = 'V. Christiaens @ UChile/ULg, Carlos Alberto Gomez Gonzalez' -__all__ = ['cube_correct_nan', - 'approx_stellar_position'] - -import copy -import numpy as np -from skimage.draw import circle -from astropy.stats import sigma_clipped_stats -from ..stats import sigma_filter - - -def cube_correct_nan(cube, neighbor_box=3, min_neighbors=3, verbose=False, - half_res_y=False): - """Sigma filtering of nan pixels in a whole frame or cube. Intended for - SINFONI data. - - Parameters - ---------- - cube : cube_like - Input 3d or 2d array. - neighbor_box : int, optional - The side of the square window around each pixel where the sigma and - median are calculated for the nan pixel correction. - min_neighbors : int, optional - Minimum number of good neighboring pixels to be able to correct the - bad/nan pixels. - verbose: {False,True} bool, optional - Whether to print more information or not during processing - half_res_y: bool, {True,False}, optional - Whether the input data have every couple of 2 rows identical, i.e. there - is twice less angular resolution vertically than horizontally (e.g. - SINFONI data). The algorithm goes twice faster if this option is - rightfully set to True. - - Returns - ------- - obj_tmp : array_like - Output cube with corrected nan pixels in each frame - """ - - obj_tmp= cube.copy() - - ndims = obj_tmp.ndim - if ndims != 2 and ndims != 3: - raise TypeError("Input object is not two or three dimensional") - - if neighbor_box < 3 or neighbor_box%2 == 0: - raise ValueError('neighbor_box should be an odd value greater than 3') - max_neigh = sum(range(3,neighbor_box+2,2)) - if min_neighbors > max_neigh: - min_neighbors = max_neigh - msg = "Warning! min_neighbors was reduced to "+str(max_neigh) +\ - " to avoid bugs. \n" - print(msg) - - def nan_corr_2d(obj_tmp): - n_x = obj_tmp.shape[1] - n_y = obj_tmp.shape[0] - - if half_res_y: - if n_y%2 != 0: - msg = 'The input frames do not have an even number of rows. ' - msg2 = 'Hence, you should probably not be using the option ' - msg3 = 'half_res_y = True.' - raise ValueError(msg+msg2+msg3) - n_y = int(n_y/2) - frame = obj_tmp - obj_tmp = np.zeros([n_y,n_x]) - for yy in range(n_y): - obj_tmp[yy] = frame[2*yy] - - # tuple with the 2D indices of each nan value of the frame - nan_indices = np.where(np.isnan(obj_tmp)) - nan_map = np.zeros_like(obj_tmp) - nan_map[nan_indices] = 1 - nnanpix = int(np.sum(nan_map)) - # Correct nan with iterative sigma filter - obj_tmp = sigma_filter(obj_tmp, nan_map, neighbor_box=neighbor_box, - min_neighbors=min_neighbors, verbose=verbose) - if half_res_y: - frame = obj_tmp - n_y = 2*n_y - obj_tmp = np.zeros([n_y,n_x]) - for yy in range(n_y): - obj_tmp[yy] = frame[int(yy/2)] - - return obj_tmp, nnanpix - - - if ndims == 2: - obj_tmp, nnanpix = nan_corr_2d(obj_tmp) - if verbose: - print('\n There were ', nnanpix, ' nan pixels corrected.') - - elif ndims == 3: - n_z = obj_tmp.shape[0] - for zz in range(n_z): - obj_tmp[zz], nnanpix = nan_corr_2d(obj_tmp[zz]) - if verbose: - msg = 'In channel '+str(zz)+', there were '+str(nnanpix) - msg2 = ' nan pixels corrected.' - print(msg+msg2) - - if verbose: - print('All nan pixels are corrected.') - - return obj_tmp - - -def approx_stellar_position(cube, fwhm, return_test=False, verbose=False): - """FIND THE APPROX COORDS OF THE STAR IN EACH CHANNEL (even the ones - dominated by noise) - - Parameters - ---------- - obj_tmp : array_like - Input 3d cube - fwhm : float or array 1D - Input full width half maximum value of the PSF for each channel. - This will be used as the standard deviation for Gaussian kernel - of the Gaussian filtering. - If float, it is assumed the same for all channels. - return_test: bool, {False,True}, optional - Whether the test result vector (a bool vector with whether the star - centroid could be find in the corresponding channel) should be returned - as well, along with the approx stellar coordinates. - verbose: {True, False}, bool optional - Chooses whether to print some additional information. - - Returns: - -------- - Array of y and x approx coordinates of the star in each channel of the cube - if return_test: it also returns the test result vector - """ - from ..phot import peak_coordinates - - obj_tmp = cube.copy() - n_z = obj_tmp.shape[0] - - if isinstance(fwhm,float) or isinstance(fwhm,int): - fwhm_scal = fwhm - fwhm = np.zeros((n_z)) - fwhm[:] = fwhm_scal - - #1/ Write a 2-columns array with indices of all max pixel values in the cube - star_tmp_idx = np.zeros([n_z,2]) - star_approx_idx = np.zeros([n_z,2]) - test_result = np.ones(n_z) - for zz in range(n_z): - star_tmp_idx[zz] = peak_coordinates(obj_tmp[zz], fwhm[zz]) - - #2/ Detect the outliers in each column - _, med_y, stddev_y = sigma_clipped_stats(star_tmp_idx[:,0],sigma=2.5) - _, med_x, stddev_x = sigma_clipped_stats(star_tmp_idx[:,1],sigma=2.5) - lim_inf_y = med_y-3*stddev_y - lim_sup_y = med_y+3*stddev_y - lim_inf_x = med_x-3*stddev_x - lim_sup_x = med_x+3*stddev_x - - if verbose: - print("median y of star - 3sigma = ", lim_inf_y) - print("median y of star + 3sigma = ", lim_sup_y) - print("median x of star - 3sigma = ", lim_inf_x) - print("median x of star + 3sigma = ", lim_sup_x) - - for zz in range(n_z): - if ((star_tmp_idx[zz,0]lim_sup_y) or - (star_tmp_idx[zz,1]lim_sup_x)): - test_result[zz] = 0 - - #3/ Replace by the median of neighbouring good coordinates if need be - for zz in range(n_z): - if test_result[zz] == 0: - ii= 1 - inf_neigh = max(0,zz-ii) - sup_neigh = min(n_z-1,zz+ii) - while test_result[inf_neigh] == 0 and test_result[sup_neigh] == 0: - ii=ii+1 - inf_neigh = max(0,zz-ii) - sup_neigh = min(n_z-1,zz+ii) - if test_result[inf_neigh] == 1 and test_result[sup_neigh] == 1: - star_approx_idx[zz] = np.floor((star_tmp_idx[sup_neigh]+ \ - star_tmp_idx[inf_neigh])/2.) - elif test_result[inf_neigh] == 1: - star_approx_idx[zz] = star_tmp_idx[inf_neigh] - else: star_approx_idx[zz] = star_tmp_idx[sup_neigh] - else: star_approx_idx[zz] = star_tmp_idx[zz] - - if return_test: - return star_approx_idx, test_result.astype(bool) - else: - return star_approx_idx diff --git a/vip_hci/preproc/derotation.py b/vip_hci/preproc/derotation.py index b8c397ba..b2a3ccc2 100644 --- a/vip_hci/preproc/derotation.py +++ b/vip_hci/preproc/derotation.py @@ -169,11 +169,10 @@ def cube_derotate(array, angle_list, imlib='opencv', interpolation='lanczos4', data_array = array pool = Pool(processes=int(nproc)) - res = pool.map(futup, itt.izip(itt.repeat(_cube_rotate_mp), - range(n_frames), itt.repeat(angle_list), - itt.repeat(imlib), - itt.repeat(interpolation), - itt.repeat(cxy))) + res = pool.map(futup, zip(itt.repeat(_cube_rotate_mp), + range(n_frames), itt.repeat(angle_list), + itt.repeat(imlib), itt.repeat(interpolation), + itt.repeat(cxy))) pool.close() array_der = np.array(res) diff --git a/vip_hci/preproc/recentering.py b/vip_hci/preproc/recentering.py index f827ecf6..d19cadbb 100644 --- a/vip_hci/preproc/recentering.py +++ b/vip_hci/preproc/recentering.py @@ -85,13 +85,13 @@ def frame_shift(array, shift_y, shift_x, imlib='opencv', image = array.copy() - if imlib=='ndimage-fourier': + if imlib == 'ndimage-fourier': shift_val = (shift_y, shift_x) array_shifted = fourier_shift(np.fft.fftn(image), shift_val) array_shifted = np.fft.ifftn(array_shifted) array_shifted = array_shifted.real - elif imlib=='ndimage-interp': + elif imlib == 'ndimage-interp': if interpolation == 'nearneig': order = 0 elif interpolation == 'bilinear': @@ -109,7 +109,7 @@ def frame_shift(array, shift_y, shift_x, imlib='opencv', array_shifted = shift(image, (shift_y, shift_x), order=order) - elif imlib=='opencv': + elif imlib == 'opencv': if no_opencv: msg = 'Opencv python bindings cannot be imported. Install opencv or ' msg += 'set imlib to ndimage-fourier or ndimage-interp' @@ -226,7 +226,8 @@ def intersection(L1, L2): raise TypeError('Input array is not a frame or 2d array') if not len(xy) == 4: raise TypeError('Input waffle spot coordinates in wrong format') - + + # TODO: verify correct handling of even/odd cases # If frame size is even we drop last row and last column if array.shape[0]%2==0: array = array[:-1,:].copy() @@ -468,11 +469,11 @@ def frame_center_radon(array, cropsize=101, hsize=0.4, step=0.01, nproc = (cpu_count()/2) pool = Pool(processes=int(nproc)) if satspots: - res = pool.map(EFT,itt.izip(itt.repeat(_radon_costf2), itt.repeat(frame), - itt.repeat(cent), itt.repeat(radint), coords)) + res = pool.map(EFT, zip(itt.repeat(_radon_costf2), itt.repeat(frame), + itt.repeat(cent), itt.repeat(radint), coords)) else: - res = pool.map(EFT,itt.izip(itt.repeat(_radon_costf), itt.repeat(frame), - itt.repeat(cent), itt.repeat(radint), coords)) + res = pool.map(EFT, zip(itt.repeat(_radon_costf), itt.repeat(frame), + itt.repeat(cent), itt.repeat(radint), coords)) costf = np.array(res) pool.close() @@ -674,7 +675,8 @@ def cube_recenter_dft_upsampling(array, cy_1, cx_1, negative=False, fwhm=4, """ if not array.ndim == 3: raise TypeError('Input array is not a cube or 3d array') - + + # TODO: verify correct handling of even/odd cases # If frame size is even we drop a row and a column if array.shape[1]%2==0: array = array[:,1:,:].copy() @@ -818,7 +820,8 @@ def cube_recenter_gauss2d_fit(array, xy, fwhm=4, subi_size=5, nproc=1, if not isinstance(pos_x,int) or not isinstance(pos_y,int): raise TypeError('pos_x and pos_y should be ints') - + + # TODO: verify correct handling of even/odd cases # If frame size is even we drop a row and a column if array.shape[1]%2==0: array = array[:,1:,:].copy() @@ -844,20 +847,15 @@ def cube_recenter_gauss2d_fit(array, xy, fwhm=4, subi_size=5, nproc=1, res = np.array(res) elif nproc>1: pool = Pool(processes=int(nproc)) - res = pool.map(EFT, itt.izip(itt.repeat(_centroid_2dg_frame), - itt.repeat(array), - range(n_frames), - subfr_sz, - itt.repeat(pos_y), - itt.repeat(pos_x), - itt.repeat(negative), - itt.repeat(debug), - fwhm, - itt.repeat(threshold))) + res = pool.map(EFT, zip(itt.repeat(_centroid_2dg_frame), + itt.repeat(array), range(n_frames), subfr_sz, + itt.repeat(pos_y), itt.repeat(pos_x), + itt.repeat(negative), itt.repeat(debug), fwhm, + itt.repeat(threshold))) res = np.array(res) pool.close() - y = cy - res[:,0] - x = cx - res[:,1] + y = cy - res[:, 0] + x = cx - res[:, 1] #return x, y if offset is not None: @@ -945,7 +943,8 @@ def cube_recenter_moffat2d_fit(array, pos_y, pos_x, fwhm=4, subi_size=5, raise TypeError('Input array is not a cube or 3d array') # if not pos_x or not pos_y: # raise ValueError('Missing parameters POS_Y and/or POS_X') - + + # TODO: verify correct handling of even/odd cases # If frame size is even we drop a row and a column if array.shape[1]%2==0: array = array[:,1:,:].copy() @@ -1002,12 +1001,11 @@ def cube_recenter_moffat2d_fit(array, pos_y, pos_x, fwhm=4, subi_size=5, res = np.array(res) elif nproc>1: pool = Pool(processes=int(nproc)) - res = pool.map(EFT,itt.izip(itt.repeat(_centroid_2dm_frame), - itt.repeat(array), range(n_frames), - size.tolist(), pos_y.tolist(), - pos_x.tolist(), star_approx_coords, - star_not_present, itt.repeat(negative), - fwhm)) + res = pool.map(EFT, zip(itt.repeat(_centroid_2dm_frame), + itt.repeat(array), range(n_frames), + size.tolist(), pos_y.tolist(), + pos_x.tolist(), star_approx_coords, + star_not_present, itt.repeat(negative), fwhm)) res = np.array(res) pool.close() y = cy - res[:,0] @@ -1193,7 +1191,8 @@ def cube_recenter_via_speckles(cube_sci, cube_ref=None, alignment_cube[0, :, :] = np.median( alignment_cube[1:(cube_sci.shape[0] + 1), :, :], axis=0) if (recenter_median): - ## Recenter the median frame using a neg. gaussian fit + # TODO: verify get_square_robust is needed. Check size=int(fwhm)+1 + # Recenter the median frame using a neg. gaussian fit sub_image, y1, x1 = get_square_robust(alignment_cube[0, :, :], size=int(fwhm) + 1, y=ceny, x=cenx, position=True) diff --git a/vip_hci/var/shapes.py b/vip_hci/var/shapes.py index d227cbef..29470900 100644 --- a/vip_hci/var/shapes.py +++ b/vip_hci/var/shapes.py @@ -80,11 +80,11 @@ def create_ringed_spider_mask(im_shape, ann_out, ann_in=0, sp_width=10, sp_angle rr0, cc0 = circle(cy, cx, min(ann_out, cy)) rr4, cc4 = circle(cy, cx, ann_in) - mask[rr0,cc0] = 1 - mask[rr1,cc1] = 0 - mask[rr2,cc2] = 0 - mask[rr3,cc3] = 0 - mask[rr4,cc4] = 0 + mask[rr0, cc0] = 1 + mask[rr1, cc1] = 0 + mask[rr2, cc2] = 0 + mask[rr3, cc3] = 0 + mask[rr4, cc4] = 0 return mask @@ -97,8 +97,8 @@ def dist(yc,xc,y1,x1): def frame_center(array, verbose=False): """ Returns the coordinates y,x of an image center. """ - cy = int(array.shape[0]/2. - 0.5) - cx = int(array.shape[1]/2. - 0.5) + cy = array.shape[0]/2. - 0.5 + cx = array.shape[1]/2. - 0.5 if verbose: print('Center px coordinates at x,y = ({:},{:})'.format(cy, cx)) @@ -112,11 +112,15 @@ def get_square(array, size, y, x, position=False): ---------- array : array_like Input frame. - size : int, odd + size : int Size of the subframe. - y, x : int - Coordinates of the center of the subframe. - position : {False, True}, optional + y : int + Y coordinate of the center of the subframe (obtained with the function + ``frame_center``). + x : int + X coordinate of the center of the subframe (obtained with the function + ``frame_center``). + position : bool optional If set to True return also the coordinates of the bottom-left vertex. Returns @@ -125,13 +129,15 @@ def get_square(array, size, y, x, position=False): Sub array. """ - if not array.ndim==2: + if not array.ndim == 2: raise TypeError('Input array is not a frame or 2d array.') - if size%2!=0: size -= 1 # making it odd to get the wing - wing = int(size/2) - # wing is added to the sides of the subframe center. Note the +1 when - # closing the interval (python doesn't include the endpoint) + # wing is added to the sides of the subframe center + if size%2 != 0: + wing = int(np.floor(size / 2.)) + else: + wing = (size / 2.) - 0.5 + # +1 because python doesn't include the endpoint when slicing array_view = array[int(y-wing):int(y+wing+1), int(x-wing):int(x+wing+1)].copy() @@ -246,15 +252,18 @@ def get_square_robust(array, size, y, x, position=False, " exceeds the borders of the array." print(msg) - if return_wings: return wing_bef,wing_aft + if return_wings: + return wing_bef,wing_aft else: # wing is added to the sides of the subframe center. Note the +1 when # closing the interval (python doesn't include the endpoint) array_view = array[int(y-wing_bef):int(y+wing_aft+1), int(x-wing_bef):int(x+wing_aft+1)].copy() - if position: return array_view, y-wing_bef, x-wing_bef - else: return array_view + if position: + return array_view, y-wing_bef, x-wing_bef + else: + return array_view def get_circle(array, radius, output_values=False, cy=None, cx=None): diff --git a/vip_hci/var/utils_var.py b/vip_hci/var/utils_var.py index 560c3d06..14d4470f 100644 --- a/vip_hci/var/utils_var.py +++ b/vip_hci/var/utils_var.py @@ -130,13 +130,18 @@ def pp_subplots(*args, **kwargs): else: show_circle = True coor_circle = kwargs['circle'] - else: show_circle = False + else: + show_circle = False - if 'circlerad' in kwargs: circle_rad = kwargs['circlerad'] - else: circle_rad = 6 + if 'circlerad' in kwargs: + circle_rad = kwargs['circlerad'] + else: + circle_rad = 6 - if 'circlealpha' in kwargs: circle_alpha = kwargs['circlealpha'] - else: circle_alpha = 0.8 + if 'circlealpha' in kwargs: + circle_alpha = kwargs['circlealpha'] + else: + circle_alpha = 0.8 # ARROW -------------------------------------------------------------------- if 'arrow' in kwargs: @@ -149,14 +154,20 @@ def pp_subplots(*args, **kwargs): else: show_arrow = False - if 'arrowshiftx' in kwargs: arrow_shiftx = kwargs['arrowshiftx'] - else: arrow_shiftx = 5 + if 'arrowshiftx' in kwargs: + arrow_shiftx = kwargs['arrowshiftx'] + else: + arrow_shiftx = 5 - if 'arrowlength' in kwargs: arrow_length = kwargs['arrowlength'] - else: arrow_length = 20 + if 'arrowlength' in kwargs: + arrow_length = kwargs['arrowlength'] + else: + arrow_length = 20 - if 'arrowalpha' in kwargs: arrow_alpha = kwargs['arrowalpha'] - else: arrow_alpha = 0.8 + if 'arrowalpha' in kwargs: + arrow_alpha = kwargs['arrowalpha'] + else: + arrow_alpha = 0.8 # LABEL -------------------------------------------------------------------- if 'label' in kwargs: @@ -167,24 +178,36 @@ def pp_subplots(*args, **kwargs): else: label = None - if 'labelsize' in kwargs: labelsize = kwargs['labelsize'] - else: labelsize = 12 + if 'labelsize' in kwargs: + labelsize = kwargs['labelsize'] + else: + labelsize = 12 - if 'labelpad' in kwargs: labelpad = kwargs['labelpad'] - else: labelpad = 5 + if 'labelpad' in kwargs: + labelpad = kwargs['labelpad'] + else: + labelpad = 5 # GRID --------------------------------------------------------------------- - if 'grid' in kwargs: grid = kwargs['grid'] - else: grid = False + if 'grid' in kwargs: + grid = kwargs['grid'] + else: + grid = False - if 'gridcolor' in kwargs: grid_color = kwargs['gridcolor'] - else: grid_color = '#f7f7f7' + if 'gridcolor' in kwargs: + grid_color = kwargs['gridcolor'] + else: + grid_color = '#f7f7f7' - if 'gridspacing' in kwargs: grid_spacing = kwargs['gridspacing'] - else: grid_spacing = 10 + if 'gridspacing' in kwargs: + grid_spacing = kwargs['gridspacing'] + else: + grid_spacing = 10 - if 'gridalpha' in kwargs: grid_alpha = kwargs['gridalpha'] - else: grid_alpha = 0.3 + if 'gridalpha' in kwargs: + grid_alpha = kwargs['gridalpha'] + else: + grid_alpha = 0.3 # VMAX-VMIN ---------------------------------------------------------------- if 'vmax' in kwargs: @@ -225,21 +248,32 @@ def pp_subplots(*args, **kwargs): else: show_cross = False if 'crossalpha' in kwargs: cross_alpha = kwargs['crossalpha'] - else: cross_alpha = 0.4 + else: + cross_alpha = 0.4 # AXIS - ANGSCALE ---------------------------------------------------------- - if 'angscale' in kwargs: angscale = kwargs['angscale'] - else: angscale = False - if 'angticksep' in kwargs: angticksep = kwargs['angticksep'] - else: angticksep = 50 - if 'pxscale' in kwargs: pxscale = kwargs['pxscale'] - else: pxscale = 0.01 # default for Keck/NIRC2 - if 'axis' in kwargs: show_axis = kwargs['axis'] - else: show_axis = True + if 'angscale' in kwargs: + angscale = kwargs['angscale'] + else: + angscale = False + if 'angticksep' in kwargs: + angticksep = kwargs['angticksep'] + else: + angticksep = 50 + if 'pxscale' in kwargs: + pxscale = kwargs['pxscale'] + else: + pxscale = 0.01 # default for Keck/NIRC2 + if 'axis' in kwargs: + show_axis = kwargs['axis'] + else: + show_axis = True # -------------------------------------------------------------------------- - if 'showcent' in kwargs: show_center = True - else: show_center = False + if 'showcent' in kwargs: + show_center = True + else: + show_center = False if 'getfig' in kwargs: getfig = kwargs['getfig'] @@ -254,7 +288,8 @@ def pp_subplots(*args, **kwargs): if 'log' in kwargs and kwargs['log'] is True: logscale = kwargs['log'] - else: logscale = False + else: + logscale = False # Defaults previously used: 'magma','CMRmap','RdBu_r' if 'cmap' in kwargs: @@ -264,27 +299,38 @@ def pp_subplots(*args, **kwargs): else: if not len(custom_cmap)==num_plots: raise RuntimeError('Cmap list does not have enough items') - else: custom_cmap = ['viridis' for i in range(num_plots)] + else: + custom_cmap = ['viridis' for i in range(num_plots)] - if 'colorb' in kwargs: colorb = kwargs['colorb'] - else: colorb = True + if 'colorb' in kwargs: + colorb = kwargs['colorb'] + else: + colorb = True - if 'dpi' in kwargs: dpi = kwargs['dpi'] - else: dpi = 90 + if 'dpi' in kwargs: + dpi = kwargs['dpi'] + else: + dpi = 90 - if 'title' in kwargs: tit = kwargs['title'] - else: tit = None + if 'title' in kwargs: + tit = kwargs['title'] + else: + tit = None - if 'horsp' in kwargs: hor_spacing = kwargs['horsp'] - else: hor_spacing = 0.3 + if 'horsp' in kwargs: + hor_spacing = kwargs['horsp'] + else: + hor_spacing = 0.3 - if 'versp' in kwargs: ver_spacing = kwargs['versp'] - else: ver_spacing = 0 + if 'versp' in kwargs: + ver_spacing = kwargs['versp'] + else: + ver_spacing = 0 # -------------------------------------------------------------------------- subplot_size = 4 - if rows==0: raise TypeError + if rows == 0: raise TypeError fig = figure(figsize=(cols*subplot_size, rows*subplot_size), dpi=dpi) if tit is not None: fig.suptitle(tit, fontsize=14) @@ -387,8 +433,6 @@ def pp_subplots(*args, **kwargs): show() - - def plot_surface(image, center=None, size=15, output=False, ds9_indexing=False, **kwargs): """