Change n_workers in parallel plotting

monash-emu · Jul 13, 2023 · f923eac · f923eac
1 parent 56b8730
commit f923eac
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/autumn/projects/sm_covid2/common_school/runner_tools.py b/autumn/projects/sm_covid2/common_school/runner_tools.py
@@ -30,6 +30,7 @@
 from autumn.projects.sm_covid2.common_school.output_plots.country_spec import make_country_output_tiling
 
 from pathlib import Path
+from multiprocessing import cpu_count
 
 countries_path = Path(PROJECTS_PATH) / "sm_covid2" / "common_school" / "included_countries.yml"
 with countries_path.open() as f:
@@ -107,11 +108,12 @@ def plot_wrapper(iso3):
         bcm = get_bcm_object(iso3, analysis)
         plot_model_fit(bcm, best_params_dict[iso3], opt_fits_path / f"best_fit_{iso3}.png")
 
-    map_parallel(plot_wrapper, iso3_list, n_workers=len(iso3_list))
+    n_workers = cpu_count()
+    map_parallel(plot_wrapper, iso3_list, n_workers=n_workers)
     if logger:
         logger.info("... finished plotting")
 
-    return best_params
+    return best_params_dict
 
 def resume_opti_search(opt, extra_iterations: int = 5000):