Merge pull request #574 from VChristiaens/master

New options for bad pixel correction + optimized NEGFC for NMF annular

vip_hci/fm/fakecomp.py

@@ -17,6 +17,10 @@
 from scipy.interpolate import interp1d
 from packaging import version
 import photutils
+try:
+    from photutils.aperture import aperture_photometry, CircularAperture
+except:
+    from photutils import aperture_photometry, CircularAperture
 if version.parse(photutils.__version__) >= version.parse('0.3'):
     # for photutils version >= '0.3' use photutils.centroids.centroid_com
     from photutils.centroids import centroid_com as cen_com
@@ -32,8 +36,8 @@
 
 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, 
+                           interpolation='lanczos4', transmission=None,
+                           radial_gradient=False, full_output=False,
                            verbose=False, nproc=1):
     """ Injects fake companions in branches, at given radial distances.
 
@@ -112,14 +116,14 @@ def cube_inject_companions(array, psf_template, angle_list, flevel, rad_dists,
     """
     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])
@@ -159,9 +163,9 @@ 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)
+                                                 -(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]))
@@ -174,19 +178,19 @@ def _cube_inject_adi(array, psf_template, angle_list, flevel, plsc,
                         fc_fr_rad = interp_trans(rad)*fc_fr
                 if nproc == 1:
                     for fr in range(nframes):
-                        array_out[fr] += _frame_shift_fcp(fc_fr_rad[fr], 
-                                                          array[fr], rad, ang, 
-                                                          angle_list[fr], 
-                                                          flevel[fr], size_fc, 
-                                                          imlib_sh, imlib_rot, 
+                        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, 
+                                                          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, 
+                                   iterable(array), rad, ang,
+                                   iterable(angle_list), iterable(flevel),
+                                   size_fc, imlib_sh, imlib_rot, interpolation,
                                    transmission, radial_gradient)
                     array_out += np.array(res)
 
@@ -318,34 +322,33 @@ 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, 
+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),
+        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)
@@ -372,9 +375,9 @@ def _frame_shift_fcp(fc_fr_rad, array, rad, ang, derot_ang, flevel, size_fc,
     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
 
 
@@ -580,7 +583,7 @@ def collapse_psf_cube(array, size, fwhm=4, verbose=True, collapse='mean'):
 
 def normalize_psf(array, fwhm='fit', size=None, threshold=None, mask_core=None,
                   model='gauss', imlib='vip-fft', interpolation='lanczos4',
-                  force_odd=True, correct_outliers=True, full_output=False, 
+                  force_odd=True, correct_outliers=True, full_output=False,
                   verbose=True, debug=False):
     """ Normalizes a PSF (2d or 3d array), to have the flux in a 1xFWHM
     aperture equal to one. It also allows to crop the array and center the PSF
@@ -656,7 +659,7 @@ def psf_norm_2d(psf, fwhm, threshold, mask_core, full_output, verbose):
                 shiftx, shifty = centrx - cx, centry - cy
                 psf = frame_shift(psf, -shifty, -shiftx, imlib=imlib,
                                   interpolation=interpolation)
-    
+
                 for _ in range(2):
                     centry, centrx = fit_2d(psf, full_output=False, debug=False)
                     if np.isnan(centry) or np.isnan(centrx):
@@ -667,9 +670,9 @@ def psf_norm_2d(psf, fwhm, threshold, mask_core, full_output, verbose):
                                       interpolation=interpolation)
 
         # we check whether the flux is normalized and fix it if needed
-        fwhm_aper = photutils.CircularAperture((cx, cy), fwhm/2)
-        fwhm_aper_phot = photutils.aperture_photometry(psf, fwhm_aper,
-                                                       method='exact')
+        fwhm_aper = CircularAperture((cx, cy), fwhm/2)
+        fwhm_aper_phot = aperture_photometry(psf, fwhm_aper,
+                                             method='exact')
         fwhm_flux = np.array(fwhm_aper_phot['aperture_sum'])
 
         if fwhm_flux > 1.1 or fwhm_flux < 0.9:
@@ -779,9 +782,9 @@ def psf_norm_2d(psf, fwhm, threshold, mask_core, full_output, verbose):
                     print_precision(fwhm)
             # Replace outliers if needed
             if correct_outliers:
-                if np.sum(np.isnan(fwhm))>0:
+                if np.sum(np.isnan(fwhm)) > 0:
                     for f in range(n):
-                        if np.isnan(fwhm[f]) and f!=0 and f!=n-1:
+                        if np.isnan(fwhm[f]) and f != 0 and f != n-1:
                             fwhm[f] = np.nanmean(np.array([fwhm[f-1],
                                                            fwhm[f+1]]))
                         elif np.isnan(fwhm[f]):

vip_hci/fm/negfc_fmerit.py

@@ -13,7 +13,7 @@
 from ..fm import cube_inject_companions
 from ..var import (frame_center, get_annular_wedge, cube_filter_highpass,
                    get_annulus_segments)
-from ..psfsub import pca_annulus, pca_annular, pca
+from ..psfsub import pca_annulus, pca_annular, nmf_annular, pca
 from ..preproc import cube_crop_frames
 
 
@@ -313,11 +313,11 @@ def get_values_optimize(cube, angs, ncomp, annulus_width, aperture_radius,
         res = pca_annulus(cube, angs, ncomp, annulus_width, r_guess, cube_ref,
                           svd_mode, scaling, imlib=imlib,
                           interpolation=interpolation, collapse=collapse,
-                          collapse_ifs=collapse_ifs, weights=weights, 
+                          collapse_ifs=collapse_ifs, weights=weights,
                           nproc=nproc)
 
-    elif algo == pca_annular:
-        
+    elif algo == pca_annular or algo == nmf_annular:
+
         tol = algo_options.get('tol', 1e-1)
         min_frames_lib = algo_options.get('min_frames_lib', 2)
         max_frames_lib = algo_options.get('max_frames_lib', 200)
@@ -333,15 +333,25 @@ def get_values_optimize(cube, angs, ncomp, annulus_width, aperture_radius,
         else:
             crop_cube = cube
             pad = 0
-        res_tmp = pca_annular(crop_cube, angs, radius_int=radius_int, fwhm=fwhm,
-                              asize=annulus_width, delta_rot=delta_rot,
-                              ncomp=ncomp, svd_mode=svd_mode, scaling=scaling,
-                              imlib=imlib, interpolation=interpolation,
-                              collapse=collapse, collapse_ifs=collapse_ifs,
-                              weights=weights, tol=tol, nproc=nproc,
-                              min_frames_lib=min_frames_lib,
-                              max_frames_lib=max_frames_lib, full_output=False,
-                              verbose=False)
+        if algo == pca_annular:
+            res_tmp = algo(crop_cube, angs, radius_int=radius_int, fwhm=fwhm,
+                           asize=annulus_width, delta_rot=delta_rot,
+                           ncomp=ncomp, svd_mode=svd_mode, scaling=scaling,
+                           imlib=imlib, interpolation=interpolation,
+                           collapse=collapse, collapse_ifs=collapse_ifs,
+                           weights=weights, tol=tol, nproc=nproc,
+                           min_frames_lib=min_frames_lib,
+                           max_frames_lib=max_frames_lib, full_output=False,
+                           verbose=False)
+        else:
+            res_tmp = algo(crop_cube, angs, radius_int=radius_int, fwhm=fwhm,
+                           asize=annulus_width, delta_rot=delta_rot,
+                           ncomp=ncomp, scaling=scaling, imlib=imlib,
+                           interpolation=interpolation, collapse=collapse,
+                           weights=weights, nproc=nproc,
+                           min_frames_lib=min_frames_lib,
+                           max_frames_lib=max_frames_lib, full_output=False,
+                           verbose=False)
         # pad again now
         res = np.pad(res_tmp, pad, mode='constant', constant_values=0)
 
@@ -358,25 +368,25 @@ def get_values_optimize(cube, angs, ncomp, annulus_width, aperture_radius,
 
     indices = disk((posy, posx), radius=aperture_radius*fwhm)
     yy, xx = indices
-    
+
     # also consider indices of the annulus for pca_annulus
     if algo == pca_annulus:
         fr_size = res.shape[-1]
         inner_rad = r_guess-annulus_width/2
         yy_a, xx_a = get_annulus_segments((fr_size, fr_size), inner_rad,
                                           annulus_width, nsegm=1)[0]
-        ## only consider overlapping indices
+        # only consider overlapping indices
         yy_f = []
         xx_f = []
         for i in range(len(yy)):
-            ind_y = np.where(yy_a==yy[i])
+            ind_y = np.where(yy_a == yy[i])
             for j in ind_y[0]:
-                if xx[i]==xx_a[j]:
+                if xx[i] == xx_a[j]:
                     yy_f.append(yy[i])
                     xx_f.append(xx[i])
         yy = np.array(yy_f, dtype=int)
         xx = np.array(xx_f, dtype=int)
-        
+
     if collapse is None:
         values = res[:, yy, xx].ravel()
     else:
@@ -527,7 +537,7 @@ def get_mu_and_sigma(cube, angs, ncomp, annulus_width, aperture_radius, fwhm,
             pad = int((cube.shape[1]-crop_sz)/2)
             crop_cube = cube_crop_frames(cube, crop_sz, verbose=False)
         else:
-            pad=0
+            pad = 0
             crop_cube = cube
 
         pca_res_tmp = pca_annular(crop_cube, angs, radius_int=radius_int,
@@ -606,4 +616,4 @@ def get_mu_and_sigma(cube, angs, ncomp, annulus_width, aperture_radius, fwhm,
     ddof = min(int(npx*(1.-(1./area)))+1, npx-1)
     sigma = np.std(all_res, ddof=ddof)
 
-    return mu, sigma
+    return mu, sigma

vip_hci/metrics/contrcurve.py

@@ -13,7 +13,10 @@
 
 import numpy as np
 import pandas as pd
-import photutils
+try:
+    from photutils.aperture import aperture_photometry, CircularAperture
+except:
+    from photutils import aperture_photometry, CircularAperture
 from inspect import getfullargspec
 from scipy.interpolate import InterpolatedUnivariateSpline
 from scipy import stats
@@ -23,16 +26,16 @@
 from ..fm import (cube_inject_companions, frame_inject_companion,
                   normalize_psf)
 from ..config import time_ini, timing
-from ..config.utils_conf import sep, vip_figsize, vip_figdpi
+from ..config.utils_conf import vip_figsize, vip_figdpi
 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, fc_rad_sep=3,
-                   noise_sep=1, wedge=(0, 360), fc_snr=100, student=True, 
-                   transmission=None, smooth=True, interp_order=2, plot=True, 
-                   dpi=vip_figdpi, debug=False, verbose=True, full_output=False, 
-                   save_plot=None, object_name=None, frame_size=None, 
+                   noise_sep=1, wedge=(0, 360), fc_snr=100, student=True,
+                   transmission=None, smooth=True, interp_order=2, plot=True,
+                   dpi=vip_figdpi, debug=False, verbose=True, full_output=False,
+                   save_plot=None, object_name=None, frame_size=None,
                    fix_y_lim=(), figsize=vip_figsize, **algo_dict):
     """ Computes the contrast curve at a given confidence (``sigma``) level for 
     an ADI cube or ADI+IFS cube. The contrast is calculated as
@@ -174,7 +177,7 @@ def contrast_curve(cube, angle_list, psf_template, fwhm, pxscale, starphot,
     if transmission is not None:
         if len(transmission) != 2 and len(transmission) != cube.shape[0]+1:
             msg = 'Wrong shape for transmission should be 2xn_rad or (nch+1) '
-            msg +='x n_rad, instead of {}'.format(transmission.shape)
+            msg += 'x n_rad, instead of {}'.format(transmission.shape)
             raise TypeError(msg)
 
     if isinstance(fwhm, (np.ndarray, list)):
@@ -204,7 +207,7 @@ def contrast_curve(cube, angle_list, psf_template, fwhm, pxscale, starphot,
         verbose_thru = True
     res_throug = throughput(cube, angle_list, psf_template, fwhm, algo=algo,
                             nbranch=nbranch, theta=theta, inner_rad=inner_rad,
-                            fc_rad_sep=fc_rad_sep, wedge=wedge, fc_snr=fc_snr, 
+                            fc_rad_sep=fc_rad_sep, wedge=wedge, fc_snr=fc_snr,
                             full_output=True, verbose=verbose_thru, **algo_dict)
     vector_radd = res_throug[3]
     if res_throug[0].shape[0] > 1:
@@ -247,7 +250,7 @@ def contrast_curve(cube, angle_list, psf_template, fwhm, pxscale, starphot,
                     ntransmission[0] = trans_rad_list
                 ntransmission[j+1] = trans_list
             transmission = ntransmission.copy()
-        if t_nz>2: #take the mean transmission over all wavelengths
+        if t_nz > 2:  # take the mean transmission over all wavelengths
             ntransmission = np.zeros([2, len(trans_rad_list)])
             ntransmission[0] = transmission[0]
             ntransmission[1] = np.mean(transmission[1:], axis=0)
@@ -657,11 +660,11 @@ def throughput(cube, angle_list, psf_template, fwhm, algo, nbranch=1, theta=0,
                                                       wedge=wedge)
     if scaling is not None:
         noise_noscal, _, _ = noise_per_annulus(frame_nofc_noscal,
-                                               separation=fwhm_med, 
+                                               separation=fwhm_med,
                                                fwhm=fwhm_med, wedge=wedge)
     else:
         noise_noscal = noise.copy()
-        
+
     vector_radd = vector_radd[inner_rad-1:]
     noise = noise[inner_rad-1:]
     res_level = res_level[inner_rad-1:]
@@ -707,7 +710,7 @@ def throughput(cube, angle_list, psf_template, fwhm, algo, nbranch=1, theta=0,
                                                      rad_dists=[radvec[i]],
                                                      theta=br*angle_branch +
                                                            theta,
-                                                     nproc=nproc, imlib=imlib, 
+                                                     nproc=nproc, imlib=imlib,
                                                      interpolation=interpolation,
                                                      verbose=False)
                     y = cy + radvec[i] * np.sin(np.deg2rad(br * angle_branch +
@@ -936,8 +939,8 @@ def find_coords(rad, sep, init_angle, fin_angle):
         yy += centery
         xx += centerx
 
-        apertures = photutils.CircularAperture(np.array((xx, yy)).T, fwhm/2)
-        fluxes = photutils.aperture_photometry(array, apertures)
+        apertures = CircularAperture(np.array((xx, yy)).T, fwhm/2)
+        fluxes = aperture_photometry(array, apertures)
         fluxes = np.array(fluxes['aperture_sum'])
 
         noise_ann = np.std(fluxes)
@@ -1003,9 +1006,9 @@ def aperture_flux(array, yc, xc, fwhm, ap_factor=1, mean=False, verbose=False):
             values = array[ind]
             obj_flux = np.mean(values)
         else:
-            aper = photutils.CircularAperture((x, y), (ap_factor*fwhm)/2)
-            obj_flux = photutils.aperture_photometry(array, aper,
-                                                     method='exact')
+            aper = CircularAperture((x, y), (ap_factor*fwhm)/2)
+            obj_flux = aperture_photometry(array, aper,
+                                           method='exact')
             obj_flux = np.array(obj_flux['aperture_sum'])
         flux[i] = obj_flux