Merge pull request #661 from VChristiaens/master

New (I)PCA features and a minor bug fix

Author: VChristiaens

vip_hci/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.6.4"
+__version__ = "1.6.5"
 
 from . import preproc
 from . import config

vip_hci/config/utils_conf.py

@@ -472,6 +472,25 @@ def pool_map(nproc, fkt, *args, **kwargs):
         if not _generator:
             res = list(res)
     else:
+        # deactivate multithreading if not yet deactivated
+        try:
+            ncpus_mt1 = os.environ["OMP_NUM_THREADS"]
+            ncpus_mt2 = os.environ["NUMEXPR_NUM_THREADS"]
+            ncpus_mt3 = os.environ["MKL_NUM_THREADS"]
+            if ncpus_mt1 != "1" or ncpus_mt2 != "1" or ncpus_mt3 != "1":
+                os.environ["OMP_NUM_THREADS"] = "1"
+                os.environ["NUMEXPR_NUM_THREADS"] = "1"
+                os.environ["MKL_NUM_THREADS"] = "1"
+                wrongly_set = True
+            else:
+                wrongly_set = False
+            vars_are_set = True
+        except KeyError:  # if the variables are not set, set them manually
+            os.environ["OMP_NUM_THREADS"] = "1"
+            os.environ["NUMEXPR_NUM_THREADS"] = "1"
+            os.environ["MKL_NUM_THREADS"] = "1"
+            vars_are_set = False
+
         # Check available start methods and pick accordingly (machine-dependent)
         avail_methods = multiprocessing.get_all_start_methods()
         # if 'forkserver' in avail_methods:  # fast and safe, if available
@@ -503,11 +522,6 @@ def pool_map(nproc, fkt, *args, **kwargs):
 
         from multiprocessing import Pool
 
-        # deactivate multithreading
-        os.environ["MKL_NUM_THREADS"] = "1"
-        os.environ["NUMEXPR_NUM_THREADS"] = "1"
-        os.environ["OMP_NUM_THREADS"] = "1"
-
         if verbose and msg is not None:
             print("{} with {} processes".format(msg, nproc))
         pool = Pool(processes=nproc)
@@ -518,11 +532,15 @@ def pool_map(nproc, fkt, *args, **kwargs):
         pool.close()
         pool.join()
 
-        # reactivate multithreading
-        ncpus = multiprocessing.cpu_count()
-        os.environ["MKL_NUM_THREADS"] = str(ncpus)
-        os.environ["NUMEXPR_NUM_THREADS"] = str(ncpus)
-        os.environ["OMP_NUM_THREADS"] = str(ncpus)
+        # return back to default behaviour regarding multithreading
+        if not vars_are_set:
+            del os.environ["OMP_NUM_THREADS"]
+            del os.environ["NUMEXPR_NUM_THREADS"]
+            del os.environ["MKL_NUM_THREADS"]
+        elif wrongly_set:
+            os.environ["OMP_NUM_THREADS"] = ncpus_mt1
+            os.environ["NUMEXPR_NUM_THREADS"] = ncpus_mt2
+            os.environ["MKL_NUM_THREADS"] = ncpus_mt3
 
     return res
 

vip_hci/fm/fakecomp.py

@@ -827,7 +827,8 @@ def cube_planet_free(planet_parameter, cube, angs, psfn, imlib='vip-fft',
         parameter must have a shape (n_pl,3) or (3,) -- the latter case assumes
         a single planet in the data. For a 4d cube r, theta and flux
         must all be 1d arrays with length equal to cube.shape[0]; i.e.
-        planet_parameter should have shape: (n_pl,3,n_ch).
+        planet_parameter should have shape: (n_pl,3,n_ch) or (3,n_ch) -- the
+        latter case assumes a single planet in the data.
     cube: numpy ndarray
         The cube of fits images expressed as a numpy.array.
     angs: numpy ndarray

vip_hci/fm/negfc_fmerit.py

@@ -621,7 +621,7 @@ def get_mu_and_sigma(cube, angs, ncomp, annulus_width, aperture_radius, fwhm,
                      cube_ref=None, wedge=None, svd_mode="lapack", scaling=None,
                      algo=pca_annulus, delta_rot=1, imlib="vip-fft",
                      interpolation="lanczos4", collapse="median", weights=None,
-                     algo_options={}, bin_spec=False):
+                     algo_options={}, bin_spec=False, verbose=False):
     """Extract the mean and standard deviation of pixel intensities in an\
     annulus of the PCA-ADI image obtained with 'algo', in the part of a defined\
     wedge that is not overlapping with PA_pl+-delta_PA.
@@ -648,7 +648,7 @@ def get_mu_and_sigma(cube, angs, ncomp, annulus_width, aperture_radius, fwhm,
         The angular position of the center of the circular aperture. This
         parameter is NOT the angular position of the candidate associated to the
         Markov chain, but should be the fixed initial guess.
-    f_guess: float, optional
+    f_guess: float or 1d numpy array, optional
         The flux estimate for the companion.
     psfn: 2D or 3D numpy ndarray, optional
         Normalized psf used to remove the companion if f_guess is provided.
@@ -721,15 +721,23 @@ def get_mu_and_sigma(cube, angs, ncomp, annulus_width, aperture_radius, fwhm,
 
     """
     if f_guess is not None and psfn is not None:
-        planet_parameter = (r_guess, theta_guess, f_guess)
+        if np.isscalar(f_guess):
+            planet_parameter = (r_guess, theta_guess, f_guess)
+        elif len(f_guess) == 1:
+            planet_parameter = (r_guess, theta_guess, f_guess[0])
+        else:
+            r_all = [r_guess]*len(f_guess)
+            theta_all = [r_guess]*len(f_guess)
+            planet_parameter = np.array([r_all, theta_all, f_guess])
         array = cube_planet_free(planet_parameter, cube, angs, psfn,
                                  imlib=imlib, interpolation=interpolation,
                                  transmission=None)
     else:
-        msg = "WARNING: f_guess not provided. The companion will not be "
-        msg += "removed from the cube before estimating mu and sigma. "
-        msg += "A wedge will be used"
-        print(msg)
+        if verbose:
+            msg = "WARNING: f_guess not provided. The companion will not be "
+            msg += "removed from the cube before estimating mu and sigma. "
+            msg += "A wedge will be used"
+            print(msg)
         array = cube.copy()
 
     centy_fr, centx_fr = frame_center(array[0])

vip_hci/fm/negfc_mcmc.py

@@ -833,7 +833,7 @@ def mcmc_negfc_sampling(cube, angs, psfn, initial_state, algo=pca_annulus,
 
     mu_sig = get_mu_and_sigma(cube, angs, ncomp, annulus_width, aperture_radius,
                               fwhm, initial_state[0], initial_state[1],
-                              initial_state[2], psfn, cube_ref=cube_ref,
+                              initial_state[2:], psfn, cube_ref=cube_ref,
                               wedge=wedge, svd_mode=svd_mode, scaling=scaling,
                               algo=algo, delta_rot=delta_rot, imlib=imlib_rot,
                               interpolation=interpolation, collapse=collapse,