added the possibility to specify the number of CPUs used by emcee

LucaMalavolta · Nov 18, 2024 · 1e9c861 · 1e9c861
1 parent f9ea4dc
commit 1e9c861
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 7 deletions.
diff --git a/SLOPpy/subroutines/bayesian_emcee.py b/SLOPpy/subroutines/bayesian_emcee.py
@@ -128,7 +128,7 @@ def emcee_lines_fit_functions(model_case,
                               lines_center,
                               jitter_index,
                               priors_dict,
-                              theta_start, boundaries, ndim, nwalkers, ngen, nsteps, nthin):
+                              theta_start, boundaries, ndim, nwalkers, ngen, nsteps, nthin, ncpus):
 
     os.environ["OMP_NUM_THREADS"] = "1"
 
@@ -173,7 +173,7 @@ def emcee_lines_fit_functions(model_case,
 
     try:
         #from pyde.de import DiffEvol
-	from pytransit.utils.de import DiffEvol
+        from pytransit.utils.de import DiffEvol
         use_pyde = True
     except ImportError:
         print('   Warnign: PyDE is not installed, random initialization point')
@@ -231,7 +231,7 @@ def emcee_lines_fit_functions(model_case,
         point_start[0, :] = theta_start
 
     start = time.time()
-    with Pool() as pool:
+    with Pool(processes=ncpus) as pool:
 
         sampler = emcee.EnsembleSampler(nwalkers,
                                         ndim,

diff --git a/SLOPpy/transmission_binned_mcmc.py b/SLOPpy/transmission_binned_mcmc.py
@@ -777,6 +777,7 @@ def compute_transmission_binned_mcmc(config_in, lines_label, reference='planetRF
         nthin = sampler_pams.get('n_thin', 50)
         nsteps = sampler_pams.get('n_steps', 20000)
         nburnin = sampler_pams.get('n_burnin', 10000)
+        ncpus = sampler_pams.get('n_cpus', 10)
         ndata = np.size(wave_meshgrid)
 
         if pams_dict.get('rp_factor', False):
@@ -809,7 +810,7 @@ def compute_transmission_binned_mcmc(config_in, lines_label, reference='planetRF
             lines_center,
             jitter_index,
             prior_dict,
-            theta_start, boundaries, ndim, nwalkers, ngen, nsteps, nthin)
+            theta_start, boundaries, ndim, nwalkers, ngen, nsteps, nthin, ncpus)
 
         flat_chain, flat_lnprob, chain_med, chain_MAP, lnprob_med, lnprob_MAP = \
             emcee_flatten_median(population, sampler_chain,
@@ -1130,6 +1131,7 @@ def compute_transmission_binned_mcmc(config_in, lines_label, reference='planetRF
         nthin = sampler_pams.get('n_thin', 50)
         nsteps = sampler_pams.get('n_steps', 20000)
         nburnin = sampler_pams.get('n_burnin', 10000)
+        nburnin = sampler_pams.get('n_cpus', 10)
         ndata = np.size(all_wave_meshgrid)
 
         if pams_dict.get('rp_factor', False):
@@ -1158,7 +1160,7 @@ def compute_transmission_binned_mcmc(config_in, lines_label, reference='planetRF
             lines_center,
             all_jitter_index,
             prior_dict,
-            theta_start, boundaries, ndim, nwalkers, ngen, nsteps, nthin)
+            theta_start, boundaries, ndim, nwalkers, ngen, nsteps, nthin, ncpus)
 
         flat_chain, flat_lnprob, chain_med, chain_MAP, lnprob_med, lnprob_MAP = \
             emcee_flatten_median(population, sampler_chain,

diff --git a/SLOPpy/transmission_mcmc.py b/SLOPpy/transmission_mcmc.py
@@ -692,6 +692,7 @@ def compute_transmission_mcmc(config_in, lines_label, reference='planetRF', pca_
         nthin = sampler_pams.get('n_thin', 50)
         nsteps = sampler_pams.get('n_steps', 20000)
         nburnin = sampler_pams.get('n_burnin', 5000)
+        ncpus = sampler_pams.get('n_cpus', 10)
         ndata = np.size(wave_array)
 
         if pams_dict.get('rp_factor', False):
@@ -720,7 +721,7 @@ def compute_transmission_mcmc(config_in, lines_label, reference='planetRF', pca_
             lines_center,
             jitter_index,
             prior_dict,
-            theta_start, boundaries, ndim, nwalkers, ngen, nsteps, nthin)
+            theta_start, boundaries, ndim, nwalkers, ngen, nsteps, nthin, ncpus)
 
         flat_chain, flat_lnprob, chain_med, chain_MAP, lnprob_med, lnprob_MAP = \
             emcee_flatten_median(population, sampler_chain,
@@ -1006,7 +1007,7 @@ def compute_transmission_mcmc(config_in, lines_label, reference='planetRF', pca_
             lines_center,
             all_jitter_index,
             prior_dict,
-            theta_start, boundaries, ndim, nwalkers, ngen, nsteps, nthin)
+            theta_start, boundaries, ndim, nwalkers, ngen, nsteps, nthin, ncpus)
 
         flat_chain, flat_lnprob, chain_med, chain_MAP, lnprob_med, lnprob_MAP = \
             emcee_flatten_median(population, sampler_chain,