diff --git a/conda-recipe/meta.yaml b/conda-recipe/meta.yaml index bb5069ba69..cb93a7be0b 100644 --- a/conda-recipe/meta.yaml +++ b/conda-recipe/meta.yaml @@ -40,6 +40,7 @@ requirements: - eko >=0.14.2 - fiatlux - sphinx >=5.0.2 + - joblib - sphinx_rtd_theme >0.5 - sphinxcontrib-bibtex - ruamel.yaml <0.18 diff --git a/n3fit/src/evolven3fit/cli.py b/n3fit/src/evolven3fit/cli.py index 3943ba6319..84f07e4ce9 100644 --- a/n3fit/src/evolven3fit/cli.py +++ b/n3fit/src/evolven3fit/cli.py @@ -2,7 +2,7 @@ def cli_evolven3fit( - configuration_folder, q_fin, q_points, op_card_info, theory_card_info, force, load, dump + configuration_folder, q_fin, q_points, op_card_info, theory_card_info, force, load, dump, ncores ): """Evolves the fitted PDFs. @@ -23,5 +23,13 @@ def cli_evolven3fit( """ utils.check_is_a_fit(configuration_folder) return evolve.evolve_fit( - configuration_folder, q_fin, q_points, op_card_info, theory_card_info, force, load, dump + configuration_folder, + q_fin, + q_points, + op_card_info, + theory_card_info, + force, + load, + dump, + ncores, ) diff --git a/n3fit/src/evolven3fit/evolve.py b/n3fit/src/evolven3fit/evolve.py index 087d3fcd70..eb6e27960d 100644 --- a/n3fit/src/evolven3fit/evolve.py +++ b/n3fit/src/evolven3fit/evolve.py @@ -5,7 +5,9 @@ import sys from ekobox import apply, genpdf, info_file +from joblib import Parallel, delayed import numpy as np +import psutil import eko from eko import basis_rotation, runner @@ -22,9 +24,19 @@ "level": logging.DEBUG, } +NUM_CORES = psutil.cpu_count(logical=False) + def evolve_fit( - fit_folder, q_fin, q_points, op_card_dict, theory_card_dict, force, eko_path, dump_eko=None + fit_folder, + q_fin, + q_points, + op_card_dict, + theory_card_dict, + force, + eko_path, + dump_eko=None, + ncores=1, ): """ Evolves all the fitted replica in fit_folder/nnfit @@ -117,10 +129,16 @@ def evolve_fit( info["NumFlavors"] = theory.heavy.num_flavs_max_pdf dump_info_file(usr_path, info) - for replica, pdf_data in initial_PDFs_dict.items(): - evolved_blocks = evolve_exportgrid(pdf_data, eko_op, x_grid) + def _wrap_evolve(pdf, replica): + evolved_blocks = evolve_exportgrid(pdf, eko_op, x_grid) dump_evolved_replica(evolved_blocks, usr_path, int(replica.removeprefix("replica_"))) + # Choose the number of cores to be the Minimal value + nb_cores = min(NUM_CORES, abs(ncores)) + Parallel(n_jobs=nb_cores)( + delayed(_wrap_evolve)(pdf, r) for r, pdf in initial_PDFs_dict.items() + ) + # remove folder: # The function dump_evolved_replica dumps the replica files in a temporary folder # We need then to remove it after fixing the position of those replica files diff --git a/n3fit/src/n3fit/scripts/evolven3fit.py b/n3fit/src/n3fit/scripts/evolven3fit.py index e37d59ef7a..c9b506b4e5 100644 --- a/n3fit/src/n3fit/scripts/evolven3fit.py +++ b/n3fit/src/n3fit/scripts/evolven3fit.py @@ -7,10 +7,10 @@ import pathlib import sys -from evolven3fit import cli, eko_utils, evolve, utils import numpy as np from eko.runner.managed import solve +from evolven3fit import cli, eko_utils, evolve, utils from n3fit.io.writer import XGRID from validphys.loader import FallbackLoader, Loader @@ -169,6 +169,7 @@ def main(): args.force, eko_path, None, + args.n_cores, ) else: # If we are in the business of producing an eko, do some checks before starting: diff --git a/pyproject.toml b/pyproject.toml index baf60ae30a..3b659e35a5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,6 +73,7 @@ reportengine = { git = "https://github.com/NNPDF/reportengine" } psutil = "*" tensorflow = "*" eko = "^0.14.1" +joblib = "*" # Hyperopt hyperopt = "*" seaborn = "*"