nnpdf 4.1.0__py3-none-any.whl → 4.1.1__py3-none-any.whl

n3fit/backends/keras_backend/MetaModel.py

@@ -1,14 +1,14 @@
 """
-    MetaModel class
+MetaModel class
 
-    Extension of the backend Model class containing some wrappers in order to absorb other
-    backend-dependent calls.
+Extension of the backend Model class containing some wrappers in order to absorb other
+backend-dependent calls.
 """
 
 from pathlib import Path
 import re
 
-from keras import Variable
+from keras import backend as K
 from keras import optimizers as Kopt
 from keras.models import Model
 import numpy as np
@@ -32,6 +32,16 @@ NN_PREFIX = "NN"
 NN_LAYER_ALL_REPLICAS = "all_NNs"
 PREPROCESSING_LAYER_ALL_REPLICAS = "preprocessing_factor"
 
+#  Running many steps in epoch eliminates some per-epoch overhead and has a big impact
+#  in GPU. In benchmarks, more than 100 steps doesn't seem to have any impact
+#  so this is the rationale for that number.
+#
+#  For reasons that are not clear at the time of writing (13/08/2025) jax only accepts
+#  one step per epoch, showing the same penalty of other libraries.
+STEPS_PER_EPOCH = 100
+if K.backend() == "jax":
+    STEPS_PER_EPOCH = 1
+
 # Some keys need to work for everyone
 for k, v in optimizers.items():
     v[1]["clipnorm"] = 1.0
@@ -98,7 +108,6 @@ class MetaModel(Model):
                 self.required_slots.add(k)
         super().__init__(input_tensors, output_tensors, **kwargs)
 
-        self.x_in = x_in
         self.input_tensors = input_tensors
         self.single_replica_generator = None
 
@@ -106,6 +115,10 @@ class MetaModel(Model):
         self.compute_losses_function = None
         self._scaler = scaler
 
+        # Keras' __setattr__ would try to track the input dictionary as a TrackedDict
+        # which is incompatible with jax, to avoid this problem, set the attribute directly
+        object.__setattr__(self, "x_in", x_in)
+
     def _parse_input(self, extra_input=None):
         """Returns the input data the model was compiled with.
         Introduces the extra_input in the places asigned to the placeholders.
@@ -153,34 +166,22 @@ class MetaModel(Model):
         if y is None:
             y = self.target_tensors
 
-        # Avoids Tensorflow overhead that happens at every epoch, by putting multiple steps in an epoch
-        steps_per_epoch = self._determine_steps_per_epoch(epochs)
+        # Running more than 1 step for every epoch eliminates some overhead of the backend libraries.
+        # In the special case in which epochs < STEPS_PER_EPOCH, set it to 1
+        if epochs < STEPS_PER_EPOCH:
+            steps_per_epoch = 1
+        else:
+            steps_per_epoch = STEPS_PER_EPOCH
 
         for k, v in x_params.items():
             x_params[k] = ops.repeat(v, steps_per_epoch, axis=0)
         y = [ops.repeat(yi, steps_per_epoch, axis=0) for yi in y]
-
         history = super().fit(
             x=x_params, y=y, epochs=epochs // steps_per_epoch, batch_size=1, **kwargs
         )
         loss_dict = history.history
         return loss_dict
 
-    def _determine_steps_per_epoch(self, epochs):
-        """Determine how many step to run in every epoch.
-        When running a single replica (CPU) or when the number of epochs is < 100 default to 1.
-        Otherwise run 100 steps per epoch.
-
-        If the number of epochs requested is not divisible by 100 there will be a number
-        of extra training epochs being run equal to max_epochs % 100 in the worst case.
-
-        """
-        num_replicas = self.output_shape[0]
-        if num_replicas == 1 or epochs < 100:
-            return 1
-
-        return 100
-
     def predict(self, x=None, **kwargs):
         """Call super().predict with the right input arguments"""
         x = self._parse_input(x)
@@ -408,7 +409,7 @@ class MetaModel(Model):
             raise ValueError("Trying to generate single replica models with no generator set.")
         replicas = []
         for i_replica in range(self.num_replicas):
-            replica = self.single_replica_generator()
+            replica = self.single_replica_generator(i_replica)
             replica.set_replica_weights(self.get_replica_weights(i_replica))
             replicas.append(replica)
 
@@ -496,9 +497,9 @@ def get_layer_replica_weights(layer, i_replica: int):
     """
     if is_stacked_single_replicas(layer):
         weights_ref = layer.get_layer(f"{NN_PREFIX}_{i_replica}").weights
-        weights = [Variable(w, name=w.name) for w in weights_ref]
+        weights = [ops.variable_to_numpy(w) for w in weights_ref]
     else:
-        weights = [Variable(w[i_replica : i_replica + 1], name=w.name) for w in layer.weights]
+        weights = [ops.variable_to_numpy(w)[i_replica : i_replica + 1] for w in layer.weights]
 
     return weights
 

n3fit/backends/keras_backend/callbacks.py

@@ -1,20 +1,21 @@
 """
-    Callbacks to be used during training
+Callbacks to be used during training
 
-    The callbacks defined in this module can be passed to the ``callbacks`` argument
-    of the ``perform_fit`` method as a list.
+The callbacks defined in this module can be passed to the ``callbacks`` argument
+of the ``perform_fit`` method as a list.
 
-    For the most typical usage: ``on_batch_end``,
-    they must take as input an epoch number and a log of the partial losses.
+For the most typical usage: ``on_batch_end``,
+they must take as input an epoch number and a log of the partial losses.
 
-    Note: the terminology used everywhere refers to a single training step as a single epoch.
-    It turns out that to avoid tensorflow overhead, it is beneficial to write a step as a
-    single batch instead. So callbacks must use ``on_batch_end``.
+Note: the terminology used everywhere refers to a single training step as a single epoch.
+It turns out that to avoid tensorflow overhead, it is beneficial to write a step as a
+single batch instead. So callbacks must use ``on_batch_end``.
 """
 
 import logging
 from time import time
 
+from keras import backend as K
 from keras.callbacks import Callback, TensorBoard
 import numpy as np
 
@@ -130,6 +131,13 @@ class StoppingCallback(CallbackStep):
         print_stats = ((epoch + 1) % self.log_freq) == 0
         # Note that the input logs correspond to the fit before the weights are updated
         logs = self.correct_logs(logs)
+
+        # WARNING: this line seems to be necessary for jax
+        # otherwise the validation model itself cannot run compute_losses
+        # but it needs to be run every epoch, which makes no sense
+        if K.backend() == "jax":
+            _ = self.model.compute_losses()
+
         self.stopping_object.monitor_chi2(logs, epoch, print_stats=print_stats)
         if self.stopping_object.stop_here():
             self.model.stop_training = True

n3fit/backends/keras_backend/internal_state.py

@@ -1,5 +1,5 @@
 """
-    Library of functions that modify the internal state of Keras/Tensorflow
+Library of functions that modify the internal state of Keras/Tensorflow
 """
 
 import os
@@ -21,7 +21,8 @@ import numpy as np
 log = logging.getLogger(__name__)
 
 # Prepare Keras-backend dependent functions
-if K.backend() in ("torch", "jax"):
+if (kback := K.backend()) == "torch":
+
     import torch
 
     def set_eager(flag=True):
@@ -56,6 +57,16 @@ elif K.backend() == "tensorflow":
                 "Could not set tensorflow parallelism settings from n3fit, maybe tensorflow is already initialized by a third program"
             )
 
+elif K.backend() == "jax":
+
+    import jax
+
+    def set_eager(flag=True):
+        pass
+
+    def set_threading(threads, core):
+        pass
+
 else:
     # Keras should've failed by now, if it doesn't it could be a new backend that works ootb?
     log.warning(f"Backend {K.backend()} not recognized. You are entering uncharted territory")

n3fit/backends/keras_backend/operations.py

@@ -1,29 +1,29 @@
 """
-    This module contains the list of operations that can be used within the
-    ``call`` method of the ``n3fit`` layers as well as operations that can
-    act on layers.
-
-    This includes an implementation of the NNPDF operations on fktable in the keras
-    language (with the mapping ``c_to_py_fun``) into Keras ``Lambda`` layers.
-
-    The rest of the operations in this module are divided into four categories:
-    numpy to tensor:
-        Operations that take a numpy array and return a tensorflow tensor
-    layer to layer:
-        Operations that take a layer and return another layer
-    tensor to tensor:
-        Operations that take a tensor and return a tensor
-    layer generation:
-        Instanciate a layer to be applied by the calling function
-
-    Most of the operations in this module are just aliases to the backend
-    (Keras in this case) so that, when implementing new backends, it is clear
-    which operations may need to be overwritten.
-    For a few selected operations, a more complicated wrapper to e.g., make
-    them into layers or apply some default, is included.
-
-    Note that tensor operations can also be applied to layers as the output of a layer is a tensor
-    equally operations are automatically converted to layers when used as such.
+This module contains the list of operations that can be used within the
+``call`` method of the ``n3fit`` layers as well as operations that can
+act on layers.
+
+This includes an implementation of the NNPDF operations on fktable in the keras
+language (with the mapping ``c_to_py_fun``) into Keras ``Lambda`` layers.
+
+The rest of the operations in this module are divided into four categories:
+numpy to tensor:
+    Operations that take a numpy array and return a tensorflow tensor
+layer to layer:
+    Operations that take a layer and return another layer
+tensor to tensor:
+    Operations that take a tensor and return a tensor
+layer generation:
+    Instanciate a layer to be applied by the calling function
+
+Most of the operations in this module are just aliases to the backend
+(Keras in this case) so that, when implementing new backends, it is clear
+which operations may need to be overwritten.
+For a few selected operations, a more complicated wrapper to e.g., make
+them into layers or apply some default, is included.
+
+Note that tensor operations can also be applied to layers as the output of a layer is a tensor
+equally operations are automatically converted to layers when used as such.
 """
 
 from keras import backend as K
@@ -70,12 +70,12 @@ elif K.backend() == "jax":
     decorator_compiler = lambda f: f
 elif K.backend() == "tensorflow":
     tensor_to_numpy_or_python = lambda x: x.numpy()
-    lambda ret: {k: i.numpy() for k, i in ret.items()}
     import tensorflow as tf
 
     decorator_compiler = tf.function
 
 dict_to_numpy_or_python = lambda ret: {k: tensor_to_numpy_or_python(i) for k, i in ret.items()}
+variable_to_numpy = lambda x: x.numpy()
 
 
 def as_layer(operation, op_args=None, op_kwargs=None, **kwargs):

n3fit/hyper_optimization/hyper_scan.py

@@ -17,10 +17,10 @@ import copy
 import logging
 import os
 
-import hyperopt
 from hyperopt.pyll.base import scope
 import numpy as np
 
+import hyperopt
 from n3fit.backends import MetaLayer, MetaModel
 from n3fit.hyper_optimization.filetrials import FileTrials
 
@@ -125,11 +125,6 @@ def hyper_scan_wrapper(replica_path_set, model_trainer, hyperscanner, max_evals=
             a ``HyperScanner`` object defining the scan
         max_evals: int
             Number of trials to run
-
-    Returns
-    -------
-        dict
-        parameters of the best trial as found by ``hyperopt``
     """
     # Tell the trainer we are doing hpyeropt
     model_trainer.set_hyperopt(True, keys=hyperscanner.hyper_keys)
@@ -181,14 +176,13 @@ def hyper_scan_wrapper(replica_path_set, model_trainer, hyperscanner, max_evals=
     )
     if hyperscanner.parallel_hyperopt:
         trials.start_mongo_workers()
-        best = hyperopt.fmin(**fmin_args, show_progressbar=True, max_queue_len=trials.num_workers)
+        hyperopt.fmin(**fmin_args, show_progressbar=True, max_queue_len=trials.num_workers)
         trials.stop_mongo_workers()
         # stop mongod command and compress database
         hyperscanner.mongod_runner.stop(mongod)
         trials.compress_mongodb_database()
     else:
-        best = hyperopt.fmin(**fmin_args, show_progressbar=False, trials_save_file=trials.pkl_file)
-    return hyperscanner.space_eval(best)
+        hyperopt.fmin(**fmin_args, show_progressbar=False, trials_save_file=trials.pkl_file)
 
 
 class ActivationStr:

n3fit/hyper_optimization/penalties.py

@@ -18,6 +18,7 @@ And return a float to be added to the hyperscan loss.
 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,12 @@ 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, ReplicaSettings
     >>> 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
+    >>> rp = [ReplicaSettings(nodes = [8], activations=["linear"], seed=0)]
+    >>> pdf_model = generate_pdf_model(rp, flav_info=fake_fl, fitbasis="FLAVOUR")
+    >>> saturation(pdf_model, 5)
+    array([0.00014878])
 
     """
     if flavors is None:
@@ -128,11 +130,12 @@ 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, ReplicaSettings
     >>> 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
+    >>> rp = [ReplicaSettings(nodes = [8], activations=["linear"], seed=0)]
+    >>> pdf_model = generate_pdf_model(rp, flav_info=fake_fl, fitbasis="FLAVOUR")
+    >>> integrability(pdf_model)
+    5.184705528587072e+21
 
     """
     pdf_instance = N3PDF(pdf_model.split_replicas())

n3fit/hyper_optimization/rewards.py

@@ -36,7 +36,7 @@ from typing import Callable
 
 import numpy as np
 
-from n3fit.vpinterface import N3PDF, compute_phi
+from n3fit.vpinterface import N3PDF, HyperoptMetrics, compute_hyperopt_metrics
 from validphys.core import DataGroupSpec
 from validphys.pdfgrids import distance_grids, xplotting_grid
 
@@ -103,7 +103,7 @@ IMPLEMENTED_STATS = {
     "best_worst": _best_worst,
     "std": _std,
 }
-IMPLEMENTED_LOSSES = ["chi2", "phi2"]
+IMPLEMENTED_LOSSES = ["chi2", "phi2", "logp", "chi2p"]
 
 
 def _pdfs_to_n3pdfs(pdfs_per_fold):
@@ -156,8 +156,10 @@ class HyperLoss:
 
         self.loss_type = self._parse_loss(loss_type)
 
-        self.phi2_vector = []
-        self.chi2_matrix = []
+        self.hyper_chi2_vector = []
+        self.hyper_phi2_vector = []
+        self.hyper_logp_vector = []
+        self.exp_chi2_matrix = []
         self.penalties = {}
 
         self.reduce_over_replicas = self._parse_statistic(replica_statistic, "replica")
@@ -167,7 +169,7 @@ class HyperLoss:
         self,
         penalties: dict[str, np.ndarray],
         validation_loss: np.ndarray,
-        kfold_loss: np.ndarray,
+        experimental_loss: np.ndarray,
         pdf_object: N3PDF,
         experimental_data: list[DataGroupSpec],
         fold_idx: int = 0,
@@ -192,8 +194,6 @@ class HyperLoss:
                 List of tuples containing `validphys.core.DataGroupSpec` instances for each group data set
             fold_idx: int
                 k-fold index. Defaults to 0.
-            include_penalties: float
-                Whether to include the penalties in the returned loss value
 
         Returns
         -------
@@ -213,7 +213,7 @@ class HyperLoss:
         >>> 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)
         """
@@ -225,21 +225,38 @@ class HyperLoss:
         best_indexes = np.argsort(validation_loss, axis=0)[:num_best]
         best_validation_losses = validation_loss[best_indexes]
 
-        # calculate phi for a given k-fold using vpinterface and validphys
+        # Select the `N3PDF` models to be used to compute the hyperopt metrics. The models
+        # are selected based on the validation losses using `self._proportion`.
         pdf_object_reduced = pdf_object.select_models(best_indexes)
-        phi2_per_fold = compute_phi(pdf_object_reduced, experimental_data) ** 2
 
-        # update hyperopt metrics
-        # these are saved in the `phi2_vector` and `chi2_matrix` attributes, excluding penalties
-        self._save_hyperopt_metrics(phi2_per_fold, kfold_loss, penalties, fold_idx)
+        # Compute the different hyperopt metrics
+        hypermetics: HyperoptMetrics = compute_hyperopt_metrics(
+            n3pdf=pdf_object_reduced, experimental_data=experimental_data
+        )
+
+        # Extract & save the values of the hyperopt metrics
+        hyper_chi2_per_fold = hypermetics.chi2  # computed with PDF covmat
+        hyper_phi2_per_fold = hypermetics.phi2  # computed without PDF covmat
+        hyper_logp_per_fold = hypermetics.logp  # computed with PDF covmat
+
+        # Update hyperopt metrics history
+        self._save_hyperopt_metrics(
+            hyper_chi2_per_fold,
+            hyper_phi2_per_fold,
+            hyper_logp_per_fold,
+            experimental_loss,
+            penalties,
+            fold_idx,
+        )
 
         # Prepare the output loss, including penalties if necessary
         if self._penalties_in_loss:
             # include penalties to experimental loss
-            kfold_loss += sum(penalties.values())
-
-            # add penalties to phi in the form of a sum of per-replicas averages
-            phi2_per_fold += sum(np.mean(penalty) for penalty in penalties.values())
+            experimental_loss += sum(penalties.values())
+            # add penalties to `phi2` and `logp` in the form of a sum of per-replicas averages
+            sum_penalties = sum(np.mean(penalty) for penalty in penalties.values())
+            hyper_phi2_per_fold += sum_penalties
+            hyper_logp_per_fold += sum_penalties
 
         # define loss for hyperopt according to the chosen loss_type
         if self.loss_type == "chi2":
@@ -247,38 +264,48 @@ class HyperLoss:
 
             # Construct the final loss as a sum of:
             # 1. The validation chi2
-            # 2. The distance to 2 for the kfold chi2
+            # 2. The distance to 2 for the experimental chi2
             # In the hyperopt paper we used 80% and 10% respectively, as a proxy for:
             # "80% of the replicas should be good, but only a small % has to cover the folds"
             # Currently take reduce_proportion for a) and 1.0 - reduce_proportion for b)
             validation_loss_average = self.reduce_over_replicas(best_validation_losses)
 
-            nselect = int(np.ceil((1.0 - self._proportion) * len(kfold_loss)))
-            best_kfold_losses = np.sort(kfold_loss, axis=0)[:nselect]
-            kfold_loss_average = self.reduce_over_replicas(best_kfold_losses)
+            nselect = int(np.ceil((1.0 - self._proportion) * len(experimental_loss)))
+            best_exp_losses = np.sort(experimental_loss, axis=0)[:nselect]
+            exp_loss_average = self.reduce_over_replicas(best_exp_losses)
 
-            loss = validation_loss_average + (max(kfold_loss_average, 2.0) - 2.0)
+            loss = validation_loss_average + (max(exp_loss_average, 2.0) - 2.0)
         elif self.loss_type == "phi2":
-            loss = phi2_per_fold
+            loss = hyper_phi2_per_fold
+        elif self.loss_type == "logp":
+            loss = hyper_logp_per_fold
+        elif self.loss_type == "chi2p":
+            loss = hyper_chi2_per_fold
 
         return loss
 
     def _save_hyperopt_metrics(
         self,
-        phi2_per_fold: float,
-        chi2_per_fold: np.ndarray,
+        hyper_chi2_per_fold: float,
+        hyper_phi2_per_fold: float,
+        hyper_logp_per_fold: float,
+        exp_chi2_per_fold: np.ndarray,
         penalties: dict[str, np.ndarray],
         fold_idx: int = 0,
     ) -> None:
         """
-        Save all chi2 and phi calculated metrics per replica and per fold, including penalties.
+        Save all the calculated metrics per replica and per fold, including penalties.
 
         Parameters
         ----------
-            phi2_per_fold: float
+            hyper_chi2_per_fold: float
+                Computed chi2 for a given k-fold
+            hyper_phi2_per_fold: float
                 Computed phi2 for a given k-fold
-            chi2_per_fold: np.ndarray
-                Computed chi2 for each replica for a given k-fold
+            hyper_logp_per_fold: float
+                Computed logp for a given k-fold
+            exp_chi2_per_fold: np.ndarray
+                Computed experimental chi2 for all the replica for a given k-fold
             penalties: Dict[str, np.ndarray]
                 dictionary of all penalties with their names
             fold_idx: int
@@ -286,13 +313,17 @@ class HyperLoss:
         """
         # reset chi2 and phi arrays for every trial
         if fold_idx == 0:
-            self.phi2_vector = []
-            self.chi2_matrix = []
+            self.hyper_chi2_vector = []
+            self.hyper_phi2_vector = []
+            self.hyper_logp_vector = []
+            self.exp_chi2_matrix = []
             self.penalties = {}
 
         # populate chi2 matrix and phi vector calculated for a given k-fold
-        self.chi2_matrix.append(chi2_per_fold)
-        self.phi2_vector.append(phi2_per_fold)
+        self.hyper_chi2_vector.append(hyper_chi2_per_fold)
+        self.hyper_phi2_vector.append(hyper_phi2_per_fold)
+        self.hyper_logp_vector.append(hyper_logp_per_fold)
+        self.exp_chi2_matrix.append(exp_chi2_per_fold)
 
         # save penalties per replica for a given k-fold
         for name, values in penalties.items():
@@ -366,7 +397,7 @@ class HyperLoss:
 
         selected_statistic = IMPLEMENTED_STATS[statistic]
 
-        if self.loss_type == "chi2":
+        if self.loss_type == "chi2" or self.loss_type == "logp" or self.loss_type == "chi2p":
             return selected_statistic
         elif self.loss_type == "phi2":
             # In case of phi2, calculate the inverse of the applied statistics