clean a bit the cli

Author: scarlehoff

n3fit/src/evolven3fit/eko_utils.py

@@ -29,6 +29,11 @@
 NFREF_DEFAULT = 5
 
 
+def _eko_theory_from_nnpdf_theory(nnpdf_theory):
+    """Takes an NNPDF theory and returns an eko theory"""
+    return runcards.Legacy(nnpdf_theory, {}).new_theory
+
+
 def construct_eko_cards(
     nnpdf_theory,
     q_fin,
@@ -53,9 +58,7 @@ def construct_eko_cards(
     # eko needs a value for Qedref and for max nf alphas
     theory["Qedref"] = theory["Qref"]
     theory["MaxNfAs"] = theory["MaxNfPdf"]
-
-    # The Legacy function is able to construct a theory card for eko starting from a NNPDF theory
-    theory_card = runcards.Legacy(theory, {}).new_theory
+    theory_card = _eko_theory_from_nnpdf_theory(theory)
 
     # construct mugrid
 
@@ -125,9 +128,7 @@ def construct_eko_photon_cards(
     # Now make sure the Legacy class still gets a Qedref, which is equal to Qref
     theory["Qedref"] = theory["Qref"]
     theory["MaxNfAs"] = theory["MaxNfPdf"]
-
-    # The Legacy function is able to construct a theory card for eko starting from a NNPDF theory
-    theory_card = runcards.Legacy(theory, {}).new_theory
+    theory_card = _eko_theory_from_nnpdf_theory(theory)
 
     # The photon needs to be evolved down to Q0
     q_fin = theory["Q0"]

n3fit/src/evolven3fit/evolve.py

@@ -112,14 +112,14 @@ def evolve_fit(
         info["XMax"] = float(x_grid[-1])
         # Save the PIDs in the info file in the same order as in the evolution
         info["Flavors"] = basis_rotation.flavor_basis_pids
-        info["NumFlavors"] = 5  # TODO: Maximum number in evol
+        info.setdefault("NumFlavors", 5)
         dump_info_file(usr_path, info)
 
         # Read the information from all the sorted replicas into what eko wants
         n_replicas = len(initial_PDFs_dict)
         all_replicas = []
         for rep_idx in range(1, n_replicas + 1):
-            # swap photon postion to match eko.basis_roation.flavor_basis_pids
+            # swap photon position to match eko.basis_rotation.flavor_basis_pids
             pdfgrid = np.array(initial_PDFs_dict[f"replica_{rep_idx}"]["pdfgrid"])
             pdfgrid = np.append(pdfgrid[:, -1].reshape(x_grid.size, 1), pdfgrid[:, :-1], axis=1)
             # and divide by x

n3fit/src/n3fit/scripts/evolven3fit.py

@@ -19,6 +19,27 @@
 _logger = logging.getLogger(__name__)
 
 
+def _add_production_options(parser):
+    """Adds the production options that are shared between the normal and the QED EKOs to their parsers"""
+    parser.add_argument("theoryID", type=int, help="ID of the theory used to produce the eko")
+    parser.add_argument("dump", type=pathlib.Path, help="Path where the EKO is dumped")
+    parser.add_argument(
+        "-i", "--x-grid-ini", default=None, type=float, help="Starting point of the x-grid"
+    )
+    parser.add_argument(
+        "-p", "--x-grid-points", default=None, type=int, help="Number of points of the x-grid"
+    )
+    parser.add_argument("-n", "--n-cores", type=int, default=1, help="Number of cores to be used")
+    parser.add_argument('--use_polarized', action='store_true', help="Use polarized evolution")
+    parser.add_argument(
+        "-e",
+        "--ev-op-iterations",
+        type=int,
+        default=None,
+        help="ev_op_iterations for the EXA theory. Overrides the settings given in the theory card.",
+    )
+
+
 def construct_eko_parser(subparsers):
     parser = subparsers.add_parser(
         "produce_eko",
@@ -30,14 +51,7 @@ def construct_eko_parser(subparsers):
             provided number of points. The eko will be dumped in the provided path."""
         ),
     )
-    parser.add_argument("theoryID", type=int, help="ID of the theory used to produce the eko")
-    parser.add_argument("dump", type=pathlib.Path, help="Path where the EKO is dumped")
-    parser.add_argument(
-        "-i", "--x-grid-ini", default=None, type=float, help="Starting point of the x-grid"
-    )
-    parser.add_argument(
-        "-p", "--x-grid-points", default=None, type=int, help="Number of points of the x-grid"
-    )
+    _add_production_options(parser)
     parser.add_argument(
         "--legacy40",
         action="store_true",
@@ -56,14 +70,7 @@ def construct_eko_photon_parser(subparsers):
             provided number of points. The eko will be dumped in the provided path."""
         ),
     )
-    parser.add_argument("theoryID", type=int, help="ID of the theory used to produce the eko")
-    parser.add_argument("dump", type=pathlib.Path, help="Path where the EKO is dumped")
-    parser.add_argument(
-        "-i", "--x-grid-ini", default=None, type=float, help="Starting point of the x-grid"
-    )
-    parser.add_argument(
-        "-p", "--x-grid-points", default=None, type=int, help="Number of points of the x-grid"
-    )
+    _add_production_options(parser)
     parser.add_argument(
         "-g", "--q-gamma", default=100, type=float, help="Scale at which the photon is generated"
     )
@@ -76,7 +83,7 @@ def construct_evolven3fit_parser(subparsers):
         help="Evolves the fitted PDFs. The q_grid starts at the Q0 given by the theory but the last point is q_fin and its number of points can be specified by q_points. If a path is given for the dump option, the eko will be dumped in that path after the computation. If a path is given for the load option, the eko to be used for the evolution will be loaded from that path. The two options are mutually exclusive.",
     )
     parser.add_argument(
-        "configuration_folder", help="Path to the folder containing the (pre-DGLAP) fit result"
+        "fit_folder", help="Path to the folder containing the (pre-DGLAP) fit result"
     )
     parser.add_argument(
         "-l", "--load", type=pathlib.Path, default=None, help="Path of the EKO to be loaded"
@@ -99,32 +106,14 @@ def evolven3fit_new():
 
 
 def main():
-    parser = ArgumentParser(
-        description="evolven3fit - a script with tools to evolve PDF fits",
-        usage="""evolven3fit [-h] [-q Q_FIN] [-p Q_POINTS] [-n N_CORES] [-e EV_OP_ITERATIONS]
-        {produce_eko,produce_eko_photon,evolve} [fit folder]
-
-        Note that with the now removed apfel-based version of `evolven3fit` the syntax was
-        `evolven3fit [fit folder] [number of replicas]`. This syntax is no longer supported in the
-        eko-based version of evolven3fit.
-        """,
-    )
-    parser.add_argument('--use_polarized', action='store_true', help="Use polarized evolution")
+    parser = ArgumentParser(description="evolven3fit - a script with tools to evolve PDF fits")
     parser.add_argument(
         "-q", "--q-fin", type=float, default=None, help="Final q-value of the evolution"
     )
     parser.add_argument(
         "-p", "--q-points", type=int, default=None, help="Number of q points for the evolution"
     )
-    parser.add_argument("-n", "--n-cores", type=int, default=1, help="Number of cores to be used")
     parser.add_argument("--no-net", action="store_true", help="Emulates validphys' --no-net")
-    parser.add_argument(
-        "-e",
-        "--ev-op-iterations",
-        type=int,
-        default=None,
-        help="ev_op_iterations for the EXA theory. Overrides the settings given in the theory card.",
-    )
 
     subparsers = parser.add_subparsers(title="actions", dest="actions")
     construct_eko_parser(subparsers)
@@ -133,12 +122,7 @@ def main():
 
     args = parser.parse_args()
 
-    op_card_info = {
-        "configs": {"n_integration_cores": args.n_cores, "polarized": args.use_polarized}
-    }
-    if args.ev_op_iterations is not None:
-        op_card_info["configs"]["ev_op_iterations"] = args.ev_op_iterations
-
+    op_card_info = {}
     # Here we do not allow any modification of the theory card, for the moment.
     theory_card_info = {}
 
@@ -149,7 +133,7 @@ def main():
 
     if args.actions == "evolve":
 
-        fit_folder = pathlib.Path(args.configuration_folder)
+        fit_folder = pathlib.Path(args.fit_folder)
         if args.load is None:
             theoryID = utils.get_theoryID_from_runcard(fit_folder)
             _logger.info(f"Loading eko from theory {theoryID}")
@@ -194,6 +178,14 @@ def main():
             )
         else:
             x_grid = np.geomspace(args.x_grid_ini, 1.0, args.x_grid_points)
+
+        # Prepare the op card config
+        op_card_info = {
+            "configs": {"n_integration_cores": args.n_cores, "polarized": args.use_polarized}
+        }
+        if args.ev_op_iterations is not None:
+            op_card_info["configs"]["ev_op_iterations"] = args.ev_op_iterations
+
         if args.actions == "produce_eko":
             tcard, opcard = eko_utils.construct_eko_cards(
                 nnpdf_theory,

n3fit/src/n3fit/vpinterface.py

@@ -1,20 +1,20 @@
 """
-    n3fit interface to validphys
+n3fit interface to validphys
 
-    Example
-    -------
+Example
+-------
 
-    >>> import numpy as np
-    >>> from n3fit.vpinterface import N3PDF
-    >>> from n3fit.model_gen import pdfNN_layer_generator
-    >>> from validphys.pdfgrids import xplotting_grid
-    >>> fake_fl = [{'fl' : i, 'largex' : [0,1], 'smallx': [1,2]} for i in ['u', 'ubar', 'd', 'dbar', 'c', 'cbar', 's', 'sbar']]
-    >>> fake_x = np.linspace(1e-3,0.8,3)
-    >>> pdf_model = pdfNN_layer_generator(nodes=[8], activations=['linear'], seed=0, flav_info=fake_fl)
-    >>> n3pdf = N3PDF(pdf_model)
-    >>> res = xplotting_grid(n3pdf, 1.6, fake_x)
-    >>> res.grid_values.error_members().shape
-    (1, 8, 3)
+>>> import numpy as np
+>>> from n3fit.vpinterface import N3PDF
+>>> from n3fit.model_gen import pdfNN_layer_generator
+>>> from validphys.pdfgrids import xplotting_grid
+>>> fake_fl = [{'fl' : i, 'largex' : [0,1], 'smallx': [1,2]} for i in ['u', 'ubar', 'd', 'dbar', 'c', 'cbar', 's', 'sbar']]
+>>> fake_x = np.linspace(1e-3,0.8,3)
+>>> pdf_model = pdfNN_layer_generator(nodes=[8], activations=['linear'], seed=0, flav_info=fake_fl)
+>>> n3pdf = N3PDF(pdf_model)
+>>> res = xplotting_grid(n3pdf, 1.6, fake_x)
+>>> res.grid_values.error_members().shape
+(1, 8, 3)
 
 
 """
@@ -25,7 +25,6 @@
 import numpy as np
 import numpy.linalg as la
 
-from n3fit.backends import PREPROCESSING_LAYER_ALL_REPLICAS
 from validphys.arclength import arc_lengths, integrability_number
 from validphys.core import PDF, MCStats
 from validphys.covmats import covmat_from_systematics, sqrt_covmat
@@ -225,9 +224,13 @@ def get_preprocessing_factors(self, replica=None):
         otherwise outputs a Nx2 array where [:,0] are alphas and [:,1] betas
         """
         # If no replica is explicitly requested, get the preprocessing layer for the first model
+        # remember replicas start counting at one
         if replica is None:
             replica = 1
-        # Replicas start counting in 1 so:
+
+        # Keep the import here to avoid loading the backend when it is not necessary
+        from n3fit.backends import PREPROCESSING_LAYER_ALL_REPLICAS
+
         preprocessing_layer = self._models[replica - 1].get_layer(PREPROCESSING_LAYER_ALL_REPLICAS)
 
         alphas_and_betas = None