finalize refactoring

Author: scarlehoff

n3fit/src/n3fit/hyper_optimization/penalties.py

@@ -18,6 +18,7 @@
 New penalties can be added directly in this module.
 The name in the runcard must match the name used in this module.
 """
+
 import numpy as np
 
 from n3fit.vpinterface import N3PDF, integrability_numbers
@@ -48,11 +49,11 @@ def saturation(pdf_model=None, n=100, min_x=1e-6, max_x=1e-4, flavors=None, **_k
     Example
     -------
     >>> from n3fit.hyper_optimization.penalties import saturation
-    >>> from n3fit.model_gen import pdfNN_layer_generator
+    >>> from n3fit.model_gen import generate_pdf_model
     >>> fake_fl = [{'fl' : i, 'largex' : [0,1], 'smallx': [1,2]} for i in ['u', 'ubar', 'd', 'dbar', 'c', 'g', 's', 'sbar']]
-    >>> pdf_model = pdfNN_layer_generator(nodes=[8], activations=['linear'], seed=0, flav_info=fake_fl, fitbasis="FLAVOUR")
-    >>> isinstance(saturation(pdf_model, 5), float)
-    True
+    >>> pdf_model = generate_pdf_model(nodes=[8], activations=['linear'], seed_list=[0], flav_info=fake_fl, fitbasis="FLAVOUR")
+    >>> saturation(pdf_model, 5), float)
+    array([0.0038])
 
     """
     if flavors is None:
@@ -128,11 +129,11 @@ def integrability(pdf_model=None, **_kwargs):
     Example
     -------
     >>> from n3fit.hyper_optimization.penalties import integrability
-    >>> from n3fit.model_gen import pdfNN_layer_generator
+    >>> from n3fit.model_gen import generate_pdf_model
     >>> fake_fl = [{'fl' : i, 'largex' : [0,1], 'smallx': [1,2]} for i in ['u', 'ubar', 'd', 'dbar', 'c', 'g', 's', 'sbar']]
-    >>> pdf_model = pdfNN_layer_generator(nodes=[8], activations=['linear'], seed=0, flav_info=fake_fl, fitbasis="FLAVOUR")
-    >>> isinstance(integrability(pdf_model), float)
-    True
+    >>> pdf_model = generate_pdf_model(nodes=[8], activations=['linear'], seed_list=[0], flav_info=fake_fl, fitbasis="FLAVOUR")
+    >>> integrability(pdf_model)
+    5.184705528587072e+21
 
     """
     pdf_instance = N3PDF(pdf_model.split_replicas())

n3fit/src/n3fit/hyper_optimization/rewards.py

@@ -213,7 +213,7 @@ def compute_loss(
         >>> ds = Loader().check_dataset("NMC_NC_NOTFIXED_P_EM-SIGMARED", variant="legacy", theoryid=399, cuts="internal")
         >>> experimental_data = [Loader().check_experiment("My DataGroupSpec", [ds])]
         >>> fake_fl = [{'fl' : i, 'largex' : [0,1], 'smallx': [1,2]} for i in ['u', 'ubar', 'd', 'dbar', 'c', 'g', 's', 'sbar']]
-        >>> pdf_model = generate_pdf_model(nodes=[8], activations=['linear'], seed=0, num_replicas=2, flav_info=fake_fl, fitbasis="FLAVOUR")
+        >>> pdf_model = generate_pdf_model(nodes=[8], activations=['linear'], seed=[0,2], flav_info=fake_fl, fitbasis="FLAVOUR")
         >>> pdf = N3PDF(pdf_model.split_replicas())
         >>> loss = hyper.compute_loss(penalties, experimental_loss, pdf, experimental_data)
         """

n3fit/src/n3fit/model_gen.py

@@ -420,7 +420,7 @@ def generate_pdf_model(
         [ list of a separate architecture per replica ]
             which can be, but is not necessary, equal for all replicas
 
-        <preprocessing factors>
+        [ <preprocessing factors> ]
             postprocessing of the network output by a variation x^{alpha}*(1-x)^{beta}
 
         <normalization>
@@ -474,29 +474,18 @@ def generate_pdf_model(
             Function to apply to the input. If given the input to the model
             will be a (1, None, 2) tensor where dim [:,:,0] is scaled
             When None, instead turn the x point into a (x, log(x)) pair
-        num_replicas: int
-            How many models should be trained in parallel.
         photons: :py:class:`validphys.photon.compute.Photon`
             If given, gives the AddPhoton layer a function to compute a photon which will be added at the
             index 0 of the 14-size FK basis
             This same function will also be used to compute the MSR component for the photon
 
-
-
     Returns
     -------
         pdf_model: MetaModel
             pdf model, with `single_replica_generator` attached as an attribute
     """
-    if len(seed_list) != num_replicas:
-        # TODO: remove this error, remove the num_replicas argument
-        raise ValueError("This should not happen")
-
-    num_replicas = len(seed_list)
-
     # Separate the settings which may be different for each replica
     # from those that are guaranteed to be equal for all replicas
-
     all_replicas = []
     for seed in seed_list:
         tmp = _ReplicaSettings(
@@ -649,44 +638,44 @@ def _pdfNN_layer_generator(
        pdf_model: n3fit.backends.MetaModel
             a model f(x) = y where x is a tensor (1, xgrid, 1) and y a tensor (1, replicas, xgrid, out)
     """
-    # TODO: at the moment nothing changes, just the signature of the function
-    seed = [i.seed for i in replicas_settings]
+    all_seed = [i.seed for i in replicas_settings]
     num_replicas = len(replicas_settings)
 
     if impose_sumrule is None:
         impose_sumrule = "All"
 
-    # Process input options. There are 2 options:
-    # 1. Scale the input
-    # 2. Concatenate log(x) to the input
-    use_feature_scaling = scaler is not None
-
-    # When scaler is active we also want to do the subtraction of large x
-    # TODO: make it its own option (i.e., one could want to use this without using scaler)
-    subtract_one = use_feature_scaling
+    ## Process the input data (x grid)
+    # There a currently two options:
+    # 1. Append log(x) to the input
+    # 2. Scale the input
+    do_nothing = lambda x: x
+    model_input = {}
 
-    # Feature scaling happens before the pdf model and changes x->(scaler(x), x),
-    # so it adds an input dimension
-    pdf_input_dimensions = 2 if use_feature_scaling else 1
-    # Adding of logs happens inside, but before the NN and adds a dimension there
-    nn_input_dimensions = 1 if use_feature_scaling else 2
+    if scaler is None:  # add log(x)
+        use_feature_scaling = subtract_one = False
+        # The PDF itself receives only x
+        pdf_input_dimensions = 1
+        # But the NN will see (x, log(x))
+        nn_input_dimensions = 2
 
-    # Define the main input
-    do_nothing = lambda x: x
-    if use_feature_scaling:
-        pdf_input = Input(shape=(None, pdf_input_dimensions), batch_size=1, name="scaledx_x")
-        process_input = do_nothing
-        extract_nn_input = Lambda(lambda x: op.op_gather_keep_dims(x, 0, axis=-1), name="x_scaled")
-        extract_original = Lambda(lambda x: op.op_gather_keep_dims(x, 1, axis=-1), name="pdf_input")
-    else:  # add log(x)
         pdf_input = Input(shape=(None, pdf_input_dimensions), batch_size=1, name="pdf_input")
         process_input = Lambda(lambda x: op.concatenate([x, op.op_log(x)], axis=-1), name="x_logx")
         extract_original = do_nothing
         extract_nn_input = do_nothing
+    else:
+        use_feature_scaling = subtract_one = True
+        # The NN will only receive x
+        nn_input_dimensions = 1
+        # But the PDF itself will receive both (x, scaler(x))
+        pdf_input_dimensions = 2
 
-    model_input = {"pdf_input": pdf_input}
+        pdf_input = Input(shape=(None, pdf_input_dimensions), batch_size=1, name="scaledx_x")
+        process_input = do_nothing
+        extract_nn_input = Lambda(lambda x: op.op_gather_keep_dims(x, 0, axis=-1), name="x_scaled")
+        extract_original = Lambda(lambda x: op.op_gather_keep_dims(x, 1, axis=-1), name="pdf_input")
 
     if subtract_one:
+        # TODO: make it its own option, even though now it only activates in the scaler if above
         input_x_eq_1 = [1.0]
         if use_feature_scaling:
             input_x_eq_1 = scaler(input_x_eq_1)[0]
@@ -695,37 +684,9 @@ def _pdfNN_layer_generator(
         layer_x_eq_1 = op.numpy_to_input(np.array(input_x_eq_1).reshape(1, 1), name="x_eq_1")
         model_input["layer_x_eq_1"] = layer_x_eq_1
 
-    # the layer that multiplies the NN output by the preprocessing factor
-    apply_preprocessing_factor = Lambda(op.op_multiply, name="prefactor_times_NN")
-
-    # Photon layer
-    layer_photon = AddPhoton(photons=photons, name="add_photon")
-
-    # Basis rotation
-    basis_rotation = FlavourToEvolution(
-        flav_info=flav_info, fitbasis=fitbasis, name="pdf_evolution_basis"
-    )
-
-    # Evolution layer
-    layer_evln = FkRotation(output_dim=out, name="pdf_FK_basis")
-
-    # Normalization and sum rules
-    if impose_sumrule:
-        sumrule_layer, integrator_input = generate_msr_model_and_grid(
-            fitbasis=fitbasis, mode=impose_sumrule, scaler=scaler, replica_seeds=seed
-        )
-        model_input["xgrid_integration"] = integrator_input
-    else:
-        sumrule_layer = lambda x: x
-
-    compute_preprocessing_factor = Preprocessing(
-        flav_info=flav_info,
-        name=PREPROCESSING_LAYER_ALL_REPLICAS,
-        replica_seeds=seed,
-        large_x=not subtract_one,
-    )
+    model_input["pdf_input"] = pdf_input
 
-    # Create the actual NeuralNetwork PDF
+    ## Create the actual NeuralNetwork PDF
     # loop over the settings for all replicas and generate a list of NN per replica
     # which will be then stack together and built into a single (input -> output) MetaModel
     # all PDFs _must_ share the same input layer
@@ -741,7 +702,21 @@ def _pdfNN_layer_generator(
     nn_pdfs = Lambda(lambda nns: op.stack(nns, axis=1), name=f"stack_replicas")(list_of_nn_pdfs)
     nn_replicas = MetaModel({'NN_input': x_input}, nn_pdfs, name=NN_LAYER_ALL_REPLICAS)
 
-    # The NN subtracted by NN(1), if applicable
+    ## Preprocessing factors:
+    # the layer that multiplies the NN output by the preprocessing factor
+    # This includes
+    #       - x^{a}(1-x)^{b}
+    #       - NN(x) - N(1.0)
+    apply_preprocessing_factor = Lambda(op.op_multiply, name="prefactor_times_NN")
+
+    compute_preprocessing_factor = Preprocessing(
+        flav_info=flav_info,
+        name=PREPROCESSING_LAYER_ALL_REPLICAS,
+        replica_seeds=all_seed,
+        large_x=not subtract_one,
+    )
+
+    # The NN subtracted by NN(1), if applicable, otherwise do nothing
     def nn_subtracted(x):
         NNs_x = nn_replicas(x)
 
@@ -752,6 +727,21 @@ def nn_subtracted(x):
 
         return NNs_x
 
+    ## Unnormalized PDF
+    #   updf_r(x) = FkRotation( NN_r(input(x)) * preprocessing_layer_r(x) )
+    #       with _r: replica index
+    #       input: whatever processing is applied to the input
+    # The preprocessing_layer and weights is specific to each replica
+    # The final PDF will be in the 14 flavours evolution basis used in the FkTables
+
+    # Basis rotation
+    basis_rotation = FlavourToEvolution(
+        flav_info=flav_info, fitbasis=fitbasis, name="pdf_evolution_basis"
+    )
+
+    # Evolution layer
+    layer_evln = FkRotation(output_dim=out, name="pdf_FK_basis")
+
     def compute_unnormalized_pdf(x):
         # Preprocess the input grid
         x_nn_input = extract_nn_input(x)
@@ -767,28 +757,39 @@ def compute_unnormalized_pdf(x):
         # Apply the preprocessing factor
         pref_NNs_x = apply_preprocessing_factor([preprocessing_factors_x, NNs_x])
 
-        # Apply basis rotation if needed
+        # Transform to FK basis, this is the full evolution basis
+        # Rotate to the 9f evolution basis first before expanding up to 14f
+        # TODO: make these two steps into one
         if not basis_rotation.is_identity():
             pref_NNs_x = basis_rotation(pref_NNs_x)
-
-        # Transform to FK basis
         PDFs_unnormalized = layer_evln(pref_NNs_x)
 
         return PDFs_unnormalized
 
     PDFs_unnormalized = compute_unnormalized_pdf(pdf_input)
 
+    ## Normalization and sum rules, produces normalized PDF
+    #   pdf_r(x) = updf_r(x) * Normalization(updf_r(integration_xgrid))
+    # The normalization layer is shared across replicas (but evaluated at each replica separately)
+    #
     if impose_sumrule:
+        sumrule_layer, integrator_input = generate_msr_model_and_grid(
+            fitbasis=fitbasis, mode=impose_sumrule, scaler=scaler, replica_seeds=all_seed
+        )
+        model_input["xgrid_integration"] = integrator_input
+
+        # We need a second unnormalized PDF evaluated on the integrated grid
         PDFs_integration_grid = compute_unnormalized_pdf(integrator_input)
 
+        # Photon contribution to the sum rule
         if photons:
             # add batch and flavor dimensions
             ph_tensor = op.numpy_to_tensor(photons.integral)
             photon_integrals = op.batchit(op.batchit(ph_tensor))
         else:
             photon_integrals = op.numpy_to_tensor(np.zeros((1, num_replicas, 1)))
 
-        PDFs_normalized = sumrule_layer(
+        PDFs = sumrule_layer(
             {
                 "pdf_x": PDFs_unnormalized,
                 "pdf_xgrid_integration": PDFs_integration_grid,
@@ -797,18 +798,21 @@ def compute_unnormalized_pdf(x):
                 "photon_integral": photon_integrals,
             }
         )
-        PDFs = PDFs_normalized
     else:
         PDFs = PDFs_unnormalized
+        sumrule_layer = lambda x: x
 
+    ## Include the photon in the PDF for QED-enabled fits
+    # (by default the entry corresponding to the photon is set to 0)
     if photons:
+        layer_photon = AddPhoton(photons=photons, name="add_photon")
         PDFs = layer_photon(PDFs)
 
+    # Return a PDF without a replica axis, to extract single replicas from an ensemble
     if not replica_axis:
         PDFs = Lambda(lambda pdfs: pdfs[:, 0], name="remove_replica_axis")(PDFs)
 
-    pdf_model = MetaModel(model_input, PDFs, name=f"PDFs", scaler=scaler)
-    return pdf_model
+    return MetaModel(model_input, PDFs, name=f"PDFs", scaler=scaler)
 
 
 # TODO: is there a way of keeping sincronized the input of this function and _ReplicaSettings
@@ -899,7 +903,7 @@ def layer_generator(i_layer, nodes_out, activation):
         previous_layer = layer(previous_layer)
 
         # Add dropout if any to the second to last layer
-        if dropout_rate > 0 and layer_idx == (len(hidden_layers) - 1):
+        if dropout_rate > 0 and layer_idx == (len(hidden_layers) - 2):
             dropout_l = base_layer_selector("dropout", rate=dropout_rate)
             previous_layer = dropout_l(previous_layer)
 

n3fit/src/n3fit/model_trainer.py

@@ -712,7 +712,6 @@ def _generate_pdf(
             regularizer_args=regularizer_args,
             impose_sumrule=self.impose_sumrule,
             scaler=self._scaler,
-            num_replicas=len(self.replicas),
             photons=photons,
         )
         return pdf_model

n3fit/src/n3fit/tests/test_hyperopt.py

@@ -19,19 +19,14 @@
 from validphys.loader import Loader
 
 
-def generate_pdf(seed, num_replicas):
+def generate_pdf(seeds):
     """Generate generic pdf model."""
     fake_fl = [
         {"fl": i, "largex": [0, 1], "smallx": [1, 2]}
         for i in ["u", "ubar", "d", "dbar", "c", "g", "s", "sbar"]
     ]
     pdf_model = generate_pdf_model(
-        nodes=[8],
-        activations=["linear"],
-        seed_list=seed,
-        num_replicas=num_replicas,
-        flav_info=fake_fl,
-        fitbasis="FLAVOUR",
+        nodes=[8], activations=["linear"], seed_list=seeds, flav_info=fake_fl, fitbasis="FLAVOUR"
     )
     return pdf_model
 
@@ -62,7 +57,7 @@ def test_compute_per_fold_loss(loss_type, replica_statistic, expected_per_fold_l
     This example assumes a 2 replica calculation with 3 added penalties.
     """
     # generate 2 replica pdf model
-    pdf_model = generate_pdf(seed=[1, 2], num_replicas=2)
+    pdf_model = generate_pdf(seeds=[1, 2])
     # add 3 penalties for a 2 replica model
     penalties = {
         'saturation': np.array([0.0, 0.0]),