Integrated nuclear fits

Created by: Radonirinaunimi

A re-implementation of the integrated nuclear-PDF fits based on the refactors.

Remaining to do:

• Improve and re-organize the writer and vpinterface modules
• Accommodate the computation of integrability and arc-lengths to properly account for A \neq 1
• Properly write the predictions onto disk and modify the way evolution is applied accordingly
• Fix broken tests
• Extend the report to include nuclear predictions

Checks:

• Standard fits in which only proton-like data are included (report)
• ...

n3fit/src/n3fit/backends/keras_backend/MetaModel.py

@@ -6,11 +6,13 @@
 """
 
 import re
+
 import numpy as np
 import tensorflow as tf
-from tensorflow.keras.models import Model
 from tensorflow.keras import optimizers as Kopt
+from tensorflow.keras.models import Model
 from tensorflow.python.keras.utils import tf_utils  # pylint: disable=no-name-in-module
+
 import n3fit.backends.keras_backend.operations as op
 
 # Check the TF version to check if legacy-mode is needed (TF < 2.2)
@@ -47,13 +49,12 @@ for k, v in optimizers.items():
     v[1]["clipnorm"] = 1.0
 
 
-def _default_loss(y_true, y_pred): # pylint: disable=unused-argument
+def _default_loss(y_true, y_pred):  # pylint: disable=unused-argument
     """Default loss to be used when the model is compiled with loss = Null
     (for instance if the prediction of the model is already the loss"""
     return op.sum(y_pred)
 
 
-
 class MetaModel(Model):
     """
     The model wraps keras.Model and adds some custom behaviour. Most notably it
@@ -95,7 +96,6 @@ class MetaModel(Model):
         if not isinstance(input_values, dict):
             raise TypeError("Expecting input_values to be a dict or None")
 
-
         x_in = {}
         # Go over the inputs. If we can deduce a constant value, either because
         # it is set in input_values or because it has a tensor_content, we
@@ -117,7 +117,6 @@ class MetaModel(Model):
         self.compute_losses_function = None
         self._scaler = scaler
 
-
     @tf.autograph.experimental.do_not_convert
     def _parse_input(self, extra_input=None):
         """Returns the input data the model was compiled with.
@@ -127,9 +126,7 @@ class MetaModel(Model):
         """
         if extra_input is None:
             if self.required_slots:
-                raise ValueError(
-                    f"The following inputs must be provided: {self.required_slots}"
-                )
+                raise ValueError(f"The following inputs must be provided: {self.required_slots}")
             return self.x_in
 
         if not isinstance(extra_input, dict):
@@ -171,10 +168,33 @@ class MetaModel(Model):
         loss_dict = history.history
         return loss_dict
 
-    def predict(self, x=None, **kwargs):
-        """ Call super().predict with the right input arguments """
+    def predict(self, x=None, proton_only=False, **kwargs):
+        """
+        Call super().predict with the right input arguments.
+
+        Parameters
+        ----------
+            x: numpy.ndarray
+               input `x` points with shape (1, Nb_x_points) at which the
+               PDF predictions will be computed
+            proton_only: boolean
+                if True return only the predictions for A=1
+
+        Returns
+        -------
+            numpy.ndarray: (Nrep, Nb_x_points, Nb_A_values, Nb_flavours)
+                array containing the PDF predictions
+        """
+        # TODO: remove `proton_only` as this should not be required
+        if proton_only:
+            x['pdf_input_A'] = op.numpy_to_tensor(np.array([[1]]))
+
         x = self._parse_input(x)
         result = super().predict(x=x, **kwargs)
+
+        if proton_only:
+            result = result[:, :, 0]
+
         return result
 
     def compute_losses(self):
@@ -325,11 +345,11 @@ class MetaModel(Model):
                 w.assign(v)
 
     def apply_as_layer(self, x):
-        """ Apply the model as a layer """
+        """Apply the model as a layer"""
         all_input = {**self.tensors_in, **x}
         return all_input, super().__call__(all_input)
 
     def get_layer_re(self, regex):
-        """ Get all layers matching the given regular expression """
+        """Get all layers matching the given regular expression"""
         check = lambda x: re.match(regex, x.name)
         return list(filter(check, self.layers))