Flavbasis

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

@@ -1,8 +1,20 @@
 """
-    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, Input
@@ -13,6 +25,14 @@
 
 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 +163,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/layers/Rotation.py → n3fit/src/n3fit/layers/Rotations.py

@@ -1,21 +1,60 @@
 """
-    This module contains rotations between basis
+    This module includes rotation layers
 """
-
 from n3fit.backends import MetaLayer
 from n3fit.backends import operations as op
+from validphys import pdfbases
 
 
 class Rotation(MetaLayer):
+    """
+    Rotates the input through some user defined rotation matrix.
+    Given an input matrix M_{m,n} with an input x_{m}, returns
+    y_{n} = x_{m}M_{m,n}
+
+    Parameters
+    ----------
+        rotation_matrix: np.array
+            rotation matrix
+        axes: int or list
+            if given a number, contracts as many indices as given
+            if given a list (of tuples) contracts indices according to op.tensor_product
+    """
+
+    def __init__(self, rotation_matrix, axes=1, **kwargs):
+        self.rotation_matrix = op.numpy_to_tensor(rotation_matrix)
+        self.axes = axes
+        super().__init__(**kwargs)
+
+    def call(self, x_raw):
+        return op.tensor_product(x_raw, self.rotation_matrix, self.axes)
+
+
+class FlavourToEvolution(Rotation):
+    """
+        Rotates from the flavour basis to
+        the evolution basis.
+    """
+
+    def __init__(
+        self, flav_info, **kwargs,
+    ):
+        rotation_matrix = pdfbases.rotation(flav_info)
+        super().__init__(rotation_matrix, axes=1, **kwargs)
+
+
+class FkRotation(MetaLayer):
     """
     Applies a transformation from the dimension-8 evolution basis
     to the dimension-14 evolution basis used by the fktables.
 
     The input to this layer is a `pdf_raw` variable which is expected to have
     a shape (1,  None, 8), and it is then rotated to an output (1, None, 14)
     """
+
     # TODO: Generate a rotation matrix in the input and just do tf.tensordot in call
     # the matrix should be: (8, 14) so that we can just do tf.tensordot(pdf, rotmat, axes=1)
+    # i.e., create the matrix and inherit from the Rotation layer above
     def __init__(self, output_dim=14, **kwargs):
         self.output_dim = output_dim
         super().__init__(**kwargs, name="evolution")

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 FkRotation, FlavourToEvolution
 from n3fit.layers.x_operations import xIntegrator, xDivide
 from n3fit.layers.MSR_Normalization import MSR_Normalization
 from n3fit.layers.DIS import DIS

n3fit/src/n3fit/model_gen.py

@@ -13,7 +13,7 @@
 from n3fit.layers import DY
 from n3fit.layers import Mask
 from n3fit.layers import ObsRotation
-from n3fit.layers import Preprocessing, Rotation
+from n3fit.layers import Preprocessing, FkRotation, FlavourToEvolution
 
 from n3fit.backends import operations
 from n3fit.backends import losses
@@ -460,12 +460,16 @@ def dense_me(x):
         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)
+    layer_evln = FkRotation(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):

n3fit/src/n3fit/performfit.py

@@ -4,6 +4,7 @@
 
 # Backend-independent imports
 from collections import namedtuple
+from validphys.pdfbases import check_basis
 import sys
 import logging
 import os.path
@@ -81,9 +82,20 @@ def check_consistent_hyperscan_options(hyperopt, hyperscan, fitting):
     if hyperopt is not None and fitting["genrep"]:
         raise CheckError("During hyperoptimization we cannot generate replicas (genrep=false)")
 
+@make_argcheck
+def check_consistent_basis(fitting):
+    fitbasis = fitting["fitbasis"]
+    # Check that there are no duplicate flavours
+    flavs = [d['fl'] for d in fitting['basis']]
+    if len(set(flavs)) != len(flavs):
+        raise CheckError(f"Repeated flavour names: check basis dictionary")
+    # Check that the basis given in the runcard is one of those defined in validphys.pdfbases
+    res = check_basis(fitbasis,flavs)
+
 # Action to be called by valid phys
 # All information defining the NN should come here in the "parameters" dict
 @check_consistent_hyperscan_options
+@check_consistent_basis
 def performfit(
     fitting,
     experiments,

n3fit/src/n3fit/tests/test_layers.py

@@ -6,6 +6,8 @@
 import numpy as np
 from n3fit.backends import operations as op
 import n3fit.layers as layers
+from validphys.pdfbases import rotation
+
 
 FLAVS = 3
 XSIZE = 4
@@ -144,3 +146,28 @@ def test_DY():
     lumi_masked = lumi_perm[mask]
     reference = np.tensordot(fk, lumi_masked, axes=3)
     assert np.allclose(result, reference, THRESHOLD)
+
+
+def test_rotation():
+    # Input dictionary to build the rotation matrix using vp2 functions
+    flav_info = [
+        {"fl": "u"},
+        {"fl": "ubar"},
+        {"fl": "d"},
+        {"fl": "dbar"},
+        {"fl": "s"},
+        {"fl": "sbar"},
+        {"fl": "c"},
+        {"fl": "g"},
+    ]
+    # Apply the rotation using numpy tensordot
+    x = np.ones(8)  # Vector in the flavour basis v_i
+    x = np.expand_dims(x, axis=[0, 1])  # Give to the input the shape (1,1,8)
+    mat = rotation(flav_info)  # Rotation matrix R_ij, i=flavour, j=evolution
+    res_np = np.tensordot(x, mat, (2, 0))  # Vector in the evolution basis u_j=R_ij*vi
+
+    # Apply the rotation through the rotation layer
+    x = op.numpy_to_tensor(x)
+    rotmat = layers.FlavourToEvolution(flav_info)
+    res_layer = rotmat(x)
+    assert np.alltrue(res_np == res_layer)

validphys2/src/validphys/pdfbases.py

@@ -427,6 +427,19 @@ def from_mapping(cls, mapping, *, aliases=None, default_elements=None):
         'v': 'V', 'v3': 'V3', 'v8': 'V8', 't3': 'T3', 't8': 'T8'},
     default_elements=(r'\Sigma', 'gluon', 'V', 'V3', 'V8', 'T3', 'T8', r'c^+', ))
 
+FLAVOUR = Basis.from_mapping(
+    {
+        'u': {'u': 1},
+        'ubar': {'ubar': 1},
+        'd': {'d': 1},
+        'dbar': {'dbar': 1},
+        's': {'s': 1},
+        'sbar': {'sbar': 1},
+        'c': {'c': 1},
+        'g': {'g': 1},
+    },
+    default_elements=('u', 'ubar', 'd', 'dbar', 's', 'sbar', 'c', 'g', ))    
+
 pdg = Basis.from_mapping({
 'g/10': {'g':0.1},
 'u_{v}': {'u':1, 'ubar':-1},
@@ -436,3 +449,50 @@ def from_mapping(cls, mapping, *, aliases=None, default_elements=None):
 r'\bar{d}': {'dbar':1},
 'c': {'c':1},
 })
+
+
+def rotation(flav_info):
+    """Return a rotation matrix R_{ij} 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), where
+    i is the flavour index and j is the evolution index. 
+    The evolution basis is defined as 
+    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).reshape(8,8)
+    # Return the transpose of the matrix, to have the first index referring to flavour
+    return mat.transpose()