From 1e9c861931bbb8308a251de4bf2bd0d8ee403ccb Mon Sep 17 00:00:00 2001
From: Gaetano Scandariato <gaetano.scandariato@inaf.it>
Date: Mon, 18 Nov 2024 13:07:31 +0100
Subject: [PATCH] added the possibility to specify the number of CPUs used by
 emcee

---
 SLOPpy/subroutines/bayesian_emcee.py | 6 +++---
 SLOPpy/transmission_binned_mcmc.py   | 6 ++++--
 SLOPpy/transmission_mcmc.py          | 5 +++--
 3 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/SLOPpy/subroutines/bayesian_emcee.py b/SLOPpy/subroutines/bayesian_emcee.py
index 651d6bc..05f1bb8 100644
--- 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
index 4520cc1..28f5f77 100644
--- 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
index 3ce8d85..2559a66 100644
--- 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,