[WIP] Running n3fit on the flavour basis

n3fit/src/n3fit/ModelTrainer.py

@@ -12,7 +12,7 @@
 import numpy as np
 import n3fit.model_gen as model_gen
 import n3fit.msr as msr_constraints
-from n3fit.backends import MetaModel, clear_backend_state
+from n3fit.backends import MetaModel, clear_backend_state, operations
 from n3fit.stopping import Stopping
 
 log = logging.getLogger(__name__)
@@ -465,15 +465,18 @@ def _generate_pdf(
             regularizer_args=regularizer_args
         )
 
-        integrator_input = None
         if self.impose_sumrule:
             # Impose the sumrule
             # Inyect here momentum sum rule, effecively modifying layer_pdf
             layer_pdf, integrator_input = msr_constraints.msr_impose(
                 layers["fitbasis"], layer_pdf
             )
-            self.input_list.append(integrator_input)
+        else:
+            nx = int(2e3)
+            xgrid, dum = msr_constraints.gen_integration_input(nx)
+            integrator_input = operations.numpy_to_input(xgrid)
 
+        self.input_list.append(integrator_input)
         self.layer_pdf = layer_pdf
 
         return layers, integrator_input

n3fit/src/n3fit/backends/keras_backend/MetaLayer.py

@@ -191,3 +191,5 @@ def permute_dimensions(self, tensor, permutation, **kwargs):
         does the permutation: axis_0 -> axis_1, axis_1 -> axis_0, axis_2 -> axis_2
         """
         return K.permute_dimensions(tensor, permutation, **kwargs)
+
+

n3fit/src/n3fit/backends/keras_backend/base_layers.py

@@ -1,18 +1,36 @@
 """
-    For a layer to be used by n3fit it should be contained in the layers dictionary below
+    This module defines custom base layers to be used by the n3fit
+    Neural Network.
+    These layers can use the keras standard set of activation function
+    or implement their own.
+
+    For a layer to be used by n3fit it should be contained in the `layers` dictionary defined below.
     This dictionary has the following structure:
 
-    'name of the layer' : ( Layer_class, {dictionary of arguments: defaults} )
+        'name of the layer' : ( Layer_class, {dictionary of arguments: defaults} )
+
+    In order to add custom activation functions, they must be added to
+    the `custom_activations` dictionary with the following structure:
+
+        'name of the activation' : function
+
+    The names of the layer and the activation function are the ones to be used in the n3fit runcard.
 """
 
 from tensorflow.keras.layers import Dense, Lambda, LSTM, Dropout, Concatenate, concatenate
 from tensorflow.keras.layers import Dense as KerasDense
 from tensorflow import expand_dims
 from tensorflow.keras.regularizers import l1_l2
-
-
 from n3fit.backends import MetaLayer
 
+# Custom activation functions
+def square_activation(x):
+    """ Squares the input """
+    return x*x
+
+custom_activations = {
+        "square" : square_activation
+        }
 
 def LSTM_modified(**kwargs):
     """
@@ -143,6 +161,9 @@ def base_layer_selector(layer_name, **kwargs):
     layer_args = layer_tuple[1]
 
     for key, value in kwargs.items():
+        # Check whether the activation function is a custom one
+        if key == "activation":
+            value = custom_activations.get(value, value)
         if key in layer_args.keys():
             layer_args[key] = value
 

n3fit/src/n3fit/backends/keras_backend/operations.py

@@ -118,3 +118,54 @@ def op_log(o_tensor, **kwargs):
     Computes the logarithm of the input
     """
     return K.log(o_tensor)
+
+
+
+def op_ratio(o_list, **kwargs):
+    """
+    Take the ratio of two observables
+    """
+    if len(o_list) != 2:
+        raise ValueError(
+            "The number of observables is incorrect, operations.py:op_ratio, expected 2, received {0}".format(
+                len(o_list)
+            )
+        )
+
+    division_layer = keras_Lambda(lambda inputs: inputs[0] / inputs[1], **kwargs)
+    return division_layer(o_list)
+
+
+def op_asy(o_list, **kwargs):
+    """
+    Perform the asymmetry operation on two observables
+    """
+    if len(o_list) != 2:
+        raise ValueError(
+            "The number of observables is incorrect, operations.py:op_asy, expected 2, received {0}".format(
+                len(o_list)
+            )
+        )
+
+    subtraction = keras_subtract(o_list)
+    addition = op_add(o_list)
+    return op_ratio([subtraction, addition], **kwargs)
+
+
+def op_smn(o_list, **kwargs):
+    """
+    Normalised sum
+    """
+    if len(o_list) != 4:
+        raise ValueError(
+            "The number of observables is incorrect, operations.py:op_smn, expected 4, received {0}".format(
+                len(o_list)
+            )
+        )
+    numer = op_add(o_list[:2])
+    denom = op_add(o_list[2:])
+    return op_ratio([numer, denom], **kwargs)
+
+def m_tensor_ones_like(*args, **kwargs):
+    return K.ones_like(*args, **kwargs)
+

n3fit/src/n3fit/layers/Preprocessing.py

@@ -36,6 +36,7 @@ def __init__(
         flav_info=None,
         seed=0,
         initializer="random_uniform",
+        flavbasis=False,
         **kwargs,
     ):
         self.output_dim = output_dim
@@ -44,6 +45,7 @@ def __init__(
         self.flav_info = flav_info
         self.seed = seed
         self.initializer = initializer
+        self.flavbasis=flavbasis
         self.kernel = []
         # super(MetaLayer, self).__init__(**kwargs)
         super().__init__(**kwargs)
@@ -90,6 +92,10 @@ def build(self, input_shape):
         # Run through the whole basis
         for flav_dict in self.flav_info:
             flav_name = flav_dict["fl"]
+            # If there are antiquarks don't generate weights for them
+            if "bar" in flav_name:
+                self.flavbasis=True
+                continue
             alpha_name = f"alpha_{flav_name}"
             beta_name = f"beta_{flav_name}"
             self.generate_weight(alpha_name, "smallx", flav_dict)
@@ -99,7 +105,22 @@ def build(self, input_shape):
 
     def call(self, inputs, **kwargs):
         x = inputs
-        pdf_raw = self.concatenate(
+        if self.flavbasis:
+            pdf_raw = self.concatenate(
+            [
+                x ** (1 - self.kernel[0][0]) * (1 - x) ** self.kernel[1][0],  # u
+                x ** (1 - self.kernel[0][0]) * (1 - x) ** self.kernel[1][0],  # ubar
+                x ** (1 - self.kernel[2][0]) * (1 - x) ** self.kernel[3][0],  # d
+                x ** (1 - self.kernel[2][0]) * (1 - x) ** self.kernel[3][0],  # dbar
+                x ** (1 - self.kernel[4][0]) * (1 - x) ** self.kernel[5][0],  # s
+                x ** (1 - self.kernel[4][0]) * (1 - x) ** self.kernel[5][0],  # sbar
+                x ** (1 - self.kernel[6][0]) * (1 - x) ** self.kernel[7][0],  # c
+                x ** (1 - self.kernel[8][0]) * (1 - x) ** self.kernel[9][0],  # g
+            ],
+            axis=-1,
+        )
+        else:   
+            pdf_raw = self.concatenate(
             [
                 x ** (1 - self.kernel[0][0]) * (1 - x) ** self.kernel[1][0],  # sigma
                 x ** (1 - self.kernel[2][0]) * (1 - x) ** self.kernel[3][0],  # g
@@ -112,4 +133,4 @@ def call(self, inputs, **kwargs):
             ],
             axis=-1,
         )
-        return pdf_raw
+        return pdf_raw    

n3fit/src/n3fit/layers/__init__.py

@@ -1,5 +1,5 @@
 from n3fit.layers.Preprocessing import Preprocessing
-from n3fit.layers.Rotation import Rotation
+from n3fit.layers.rotations import Rotation, FlavourToEvolution
 from n3fit.layers.x_operations import xIntegrator, xDivide, xMultiply
 from n3fit.layers.MSR_Normalization import MSR_Normalization
 from n3fit.layers.DIS import DIS

n3fit/src/n3fit/layers/Rotation.py → n3fit/src/n3fit/layers/rotations.py

@@ -1,5 +1,26 @@
+"""
+    This module includes rotation layers
+"""
 from n3fit.backends import MetaLayer
+from validphys.pdfbases import rotation
 
+class FlavourToEvolution(MetaLayer):
+    """ 
+        Rotates from the flavour basis to
+        the evolution basis. 
+    """
+    def __init__(
+        self,
+        flav_info,
+        **kwargs,
+    ):
+        rotation_matrix = rotation(flav_info)
+        self.rotation_matrix = self.np_to_tensor(rotation_matrix)   
+        super().__init__(**kwargs)
+
+    def call(self, x_raw):     
+        return self.tensor_product(x_raw, self.rotation_matrix, 1)      
+
 
 class Rotation(MetaLayer):
     """

n3fit/src/n3fit/model_gen.py

@@ -12,8 +12,8 @@
 from n3fit.layers import DIS
 from n3fit.layers import DY
 from n3fit.layers import Mask
+from n3fit.layers import Preprocessing, Rotation, FlavourToEvolution
 from n3fit.layers import ObsRotation
-from n3fit.layers import Preprocessing, Rotation
 
 from n3fit.backends import operations
 from n3fit.backends import losses
@@ -435,28 +435,37 @@ def pdfNN_layer_generator(
     def dense_me(x):
         """ Takes an input tensor `x` and applies all layers
         from the `list_of_pdf_layers` in order """
+        x0 = operations.m_tensor_ones_like(x)
         if inp == 1:
             curr_fun = list_of_pdf_layers[0](x)
+            curr_fun0 = list_of_pdf_layers[0](x0)
         else:
             curr_fun = list_of_pdf_layers[0](add_log(x))
+            curr_fun0 = list_of_pdf_layers[0](add_log(x0))
 
         for dense_layer in list_of_pdf_layers[1:]:
             curr_fun = dense_layer(curr_fun)
-        return curr_fun
+            curr_fun0 = dense_layer(curr_fun0)
+        res = operations.op_subtract([curr_fun,curr_fun0])    
+        return res
 
     # Preprocessing layer (will be multiplied to the last of the denses)
     preproseed = seed + number_of_layers
     layer_preproc = Preprocessing(
         input_shape=(1,), name="pdf_prepro", flav_info=flav_info, seed=preproseed
     )
 
-    # Apply preprocessing
-    def layer_fitbasis(x):
-        return operations.op_multiply([dense_me(x), layer_preproc(x)])
-
     # Evolution layer
     layer_evln = Rotation(input_shape=(last_layer_nodes,), output_dim=out)
 
+    # Basis rotation
+    basis_rotation = FlavourToEvolution(flav_info=flav_info)
+
+    # Apply preprocessing and basis
+    def layer_fitbasis(x):
+        ret = operations.op_multiply([dense_me(x), layer_preproc(x)])
+        return basis_rotation(ret)
+
     # Rotation layer, changes from the 8-basis to the 14-basis
     def layer_pdf(x):
         return layer_evln(layer_fitbasis(x))

validphys2/src/validphys/pdfbases.py

@@ -394,3 +394,48 @@ def from_mapping(cls, mapping, *, aliases=None, default_elements=None):
 r'\bar{d}': {'dbar':1},
 'c': {'c':1},
 })
+
+
+def rotation(flav_info):
+    """Returnd a rotation matrix which takes from the flavour to the evolution basis,
+    from (u, ubar, d, dbar, s, sbar, c, g) to (sigma, g, v, v3, v8, t3, t8, cp) 
+    with 
+    cp = c + cbar = 2c
+    and
+    sigma = u + ubar + d + dbar + s + sbar + cp  
+    v = u - ubar + d - dbar + s - sbar + c - cbar 
+    v3 = u - ubar - d + dbar
+    v8 = u - ubar + d - dbar - 2*s + 2*sbar
+    t3 = u + ubar - d - dbar
+    t8 = u + ubar + d + dbar - 2*s - 2*sbar
+
+    If the input is already in the evolution basis it returns the identity.
+    """
+    sigma = {'u': 1, 'ubar': 1, 'd': 1, 'dbar': 1, 's': 1, 'sbar': 1, 'c': 2, 'g': 0 }
+    v = {'u': 1, 'ubar': -1, 'd': 1, 'dbar': -1, 's': 1, 'sbar': -1, 'c': 0, 'g': 0 }
+    v3 = {'u': 1, 'ubar': -1, 'd': -1, 'dbar': 1, 's': 0, 'sbar': 0, 'c': 0, 'g': 0 }
+    v8 = {'u': 1, 'ubar': -1, 'd': 1, 'dbar': -1, 's': -2, 'sbar': 2, 'c': 0, 'g': 0 }
+    t3 = {'u': 1, 'ubar': 1, 'd': -1, 'dbar': -1, 's': 0, 'sbar': 0, 'c': 0, 'g': 0 }
+    t8 = {'u': 1, 'ubar': 1, 'd': 1, 'dbar': 1, 's': -2, 'sbar': -2, 'c': 0, 'g': 0 }
+    cp = {'u': 0, 'ubar': 0, 'd': 0, 'dbar': 0, 's': 0, 'sbar': 0, 'c': 2, 'g': 0 }
+    g = {'u': 0, 'ubar': 0, 'd': 0, 'dbar': 0, 's': 0, 'sbar': 0, 'c': 0, 'g': 1 }
+    flist = [sigma, g, v, v3, v8, t3, t8, cp]
+
+    evol_basis = False
+    mat = []
+    for f in flist:
+        for flav_dict in flav_info:
+            try:
+                flav_name = flav_dict["fl"]
+                mat.append(f[flav_name])
+            # if one of the keys in the dictionary is not a key in flist
+            # it means we are already in the evolution basis    
+            except KeyError:
+                evol_basis = True
+                break
+        if evol_basis:
+            mat = np.identity(8)
+            break    
+
+    mat = np.asarray(mat)
+    return mat.reshape(8,8)