diff --git a/n3fit/src/n3fit/backends/keras_backend/base_layers.py b/n3fit/src/n3fit/backends/keras_backend/base_layers.py index 2e00b28fbf..61db66a13e 100644 --- a/n3fit/src/n3fit/backends/keras_backend/base_layers.py +++ b/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 diff --git a/n3fit/src/n3fit/layers/Rotation.py b/n3fit/src/n3fit/layers/Rotations.py similarity index 54% rename from n3fit/src/n3fit/layers/Rotation.py rename to n3fit/src/n3fit/layers/Rotations.py index 65b7b5f8af..82e5105a2e 100644 --- a/n3fit/src/n3fit/layers/Rotation.py +++ b/n3fit/src/n3fit/layers/Rotations.py @@ -1,12 +1,49 @@ """ - 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. @@ -14,8 +51,10 @@ class Rotation(MetaLayer): 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") diff --git a/n3fit/src/n3fit/layers/__init__.py b/n3fit/src/n3fit/layers/__init__.py index 0efc973649..5fae6c64c1 100644 --- a/n3fit/src/n3fit/layers/__init__.py +++ b/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 diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py index fdb0313e66..04fe443129 100644 --- a/n3fit/src/n3fit/model_gen.py +++ b/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): diff --git a/n3fit/src/n3fit/performfit.py b/n3fit/src/n3fit/performfit.py index 8b3f3834e9..3511f795ba 100644 --- a/n3fit/src/n3fit/performfit.py +++ b/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, diff --git a/n3fit/src/n3fit/tests/test_layers.py b/n3fit/src/n3fit/tests/test_layers.py index 14ad70bbc2..f9b953dfa2 100644 --- a/n3fit/src/n3fit/tests/test_layers.py +++ b/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) diff --git a/validphys2/src/validphys/pdfbases.py b/validphys2/src/validphys/pdfbases.py index 4e3e38535f..1b492def98 100644 --- a/validphys2/src/validphys/pdfbases.py +++ b/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()