Move the operations from MetaLayer to operations.py

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

@@ -114,80 +114,3 @@ def select_initializer(ini_name, seed=None, **kwargs):
             if key in ini_args.keys():
                 ini_args[key] = value
         return ini_class(**ini_args)
-
-    # Make numpy array into a tensor
-    def np_to_tensor(self, np_array, **kwargs):
-        """
-        Given a numpy array, returns a constant tensor
-        """
-        return K.constant(np_array, **kwargs)
-
-    # Common tensor operations
-    def tensor_ones(self, shape, **kwargs):
-        """
-        Generates a tensor of ones of the given shape
-        """
-        return K.ones(shape, **kwargs)
-
-    def tensor_ones_like(self, tensor, **kwargs):
-        """
-        Generates a tensor of ones of the same shape as the input tensor
-        """
-        return K.ones_like(tensor, **kwargs)
-
-    def many_replication(self, grid, replications, axis=0, **kwargs):
-        """
-            Generates a tensor with one extra dimension:
-                a repetition of "grid" n times along the given axis
-            from keras documentation:
-            If x has shape (s1, s2, s3) and axis is 1, the output will have shape (s1, s2 * rep, s3)
-        """
-        return K.repeat_elements(grid, rep=replications, axis=axis, **kwargs)
-
-    def sum(self, tensor, axis=None, **kwargs):
-        """
-        Computes the sum of the elements of the tensor
-        """
-        return K.sum(tensor, axis=axis, **kwargs)
-
-    def tensor_product(self, tensor_x, tensor_y, axes, **kwargs):
-        """
-        Computes the tensordot product between tensor_x and tensor_y
-        """
-        return tf.tensordot(tensor_x, tensor_y, axes=axes, **kwargs)
-
-    def transpose(self, tensor, **kwargs):
-        """
-        Transpose a tensor
-        """
-        return K.transpose(tensor, **kwargs)
-
-    def boolean_mask(self, tensor, mask, axis=None, **kwargs):
-        """
-        Applies boolean mask to a tensor
-        """
-        return tf.boolean_mask(tensor, mask, axis=axis, **kwargs)
-
-    def concatenate(self, tensor_list, axis=-1, target_shape=None):
-        """
-        Concatenates a list of numbers or tenosr into a bigger tensor
-        If the target shape is given, the output is reshaped to said shape
-        """
-        concatenated_tensor = K.concatenate(tensor_list, axis=axis)
-        if target_shape:
-            return K.reshape(concatenated_tensor, target_shape)
-        else:
-            return concatenated_tensor
-
-    def flatten(self, x):
-        """ Flatten tensor x """
-        return tf.reshape(x, (-1,))
-
-    def permute_dimensions(self, tensor, permutation, **kwargs):
-        """
-        Receives a tensor and a tuple and permutes the axes of the tensor according to it.
-        i.e.
-        if permutation = (1,0,2)
-        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/MetaModel.py

@@ -8,6 +8,18 @@
 from tensorflow.keras.models import Model
 from tensorflow.keras import optimizers as Kopt
 
+# Define in this dictionary new optimizers as well as the arguments they accept
+# (with default values if needed be)
+optimizers = {
+    "RMSprop": (Kopt.RMSprop, {"lr": 0.01}),
+    "Adam": (Kopt.Adam, {"lr": 0.01}),
+    "Adagrad": (Kopt.Adagrad, {}),
+    "Adadelta": (Kopt.Adadelta, {"lr": 1.0}),
+    "Adamax": (Kopt.Adamax, {}),
+    "Nadam": (Kopt.Nadam, {}),
+    "Amsgrad": (Kopt.Adam, {"lr": 0.01, "amsgrad": True}),
+}
+
 
 def _fill_placeholders(original_input, new_input=None):
     """
@@ -57,18 +69,6 @@ class MetaModel(Model):
             keyword arguments to pass directly to Model
     """
 
-    # Define in this dictionary new optimizers as well as the arguments they accept
-    # (with default values if needed be)
-    optimizers = {
-        "RMSprop": (Kopt.RMSprop, {"lr": 0.01}),
-        "Adam": (Kopt.Adam, {"lr": 0.01}),
-        "Adagrad": (Kopt.Adagrad, {}),
-        "Adadelta": (Kopt.Adadelta, {"lr": 1.0}),
-        "Adamax": (Kopt.Adamax, {}),
-        "Nadam": (Kopt.Nadam, {}),
-        "Amsgrad": (Kopt.Adam, {"lr": 0.01, "amsgrad": True}),
-    }
-
     def __init__(self, input_tensors, output_tensors, **kwargs):
         self.has_dataset = False
 
@@ -227,7 +227,7 @@ def compile(
                 if given further calls to fit/evaluate must be done with y = None.
         """
         try:
-            opt_tuple = self.optimizers[optimizer_name]
+            opt_tuple = optimizers[optimizer_name]
         except KeyError as e:
             raise NotImplementedError(
                 f"[MetaModel.select_initializer] optimizer not implemented: {optimizer_name}"

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

@@ -1,12 +1,22 @@
 """
-    This module containg a list of useful operations translated in the keras language
-
-    All operations accept as input an iterable of keras layers or tensors
-    and (when necessary) keyword arguments.
-    The return operation is always a keras layer (or tensor)
+    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 (hence the mapping `c_to_py_fun`)
+    language (with the mapping ``c_to_py_fun``) into Keras ``Lambda`` layers.
+
+    The rest of the operations in this module are divided into three categories:
+    numpy to tensor:
+        Operations that take a numpy array and return a tensorflow tensor
+    tensor to tensor:
+        Operations that take a tensor and return a tensor
+    layer generation:
+        Instanciate a layer to be applied by the calling function
+
+    Some of these are just aliases to the backend (tensorflow or Keras) operations
+    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.
 """
 
 import numpy as np
@@ -15,24 +25,53 @@
 from tensorflow.keras.layers import multiply as keras_multiply
 from tensorflow.keras.layers import Concatenate as keras_concatenate
 
-from tensorflow.keras.layers import Input, Layer
+from tensorflow.keras.layers import Input
 from tensorflow.keras import backend as K
 
 from validphys.convolution import OP
 
 
+def evaluate(tensor):
+    """ Evaluate input tensor using the backend """
+    return K.eval(tensor)
+
+
+# NNPDF operations
+def c_to_py_fun(op_name, name="dataset"):
+    """
+    Map the NNPDF operations to Keras layers
+    NNPDF operations are defined in :py:func:`validphys.convolution.OP`
+
+    Parameters
+    ----------
+        op_name: str
+            A string defining the operation name
+    """
+    try:
+        operation = OP[op_name]
+    except KeyError as e:
+        raise ValueError(f"Operation {op_name} not recognised") from e
+
+    # Convert the operation into a lambda layer
+    operation_layer = keras_Lambda(lambda x: operation(*x), name=f"op_{name}_{op_name}")
+    return operation_layer
+
+
+# f(x: numpy) -> y: tensor
 def numpy_to_tensor(ival):
     """
         Make the input into a tensor
     """
     return K.constant(ival)
 
 
+# f(x: tensor) -> y: tensor
 def batchit(x, batch_dimension=0):
     """ Add a batch dimension to tensor x """
     return tf.expand_dims(x, batch_dimension)
 
 
+# layer generation
 def concatenate_split(splitting_sizes, axis=1):
     """ Generate a pair of concatention and splitting layer
     so that they invert each other
@@ -49,6 +88,7 @@ def concatenate_split(splitting_sizes, axis=1):
     return concatenation_layer, splitting_layer
 
 
+# layer generation
 def numpy_to_input(numpy_array, no_reshape=False, name=None):
     """
     Takes a numpy array and generates a Input layer.
@@ -77,33 +117,76 @@ def numpy_to_input(numpy_array, no_reshape=False, name=None):
     return input_layer
 
 
-def evaluate(tensor):
-    """ Evaluate input tensor using the backend """
-    return K.eval(tensor)
+#
+# Tensor operations
+# f(x: tensor[s]) -> y: tensor
+#
+
+# Generation operations
+# generate tensors of given shape/content
+def tensor_ones_like(*args, **kwargs):
+    """
+    Generates a tensor of ones of the same shape as the input tensor
+    See full `docs <https://www.tensorflow.org/api_docs/python/tf/keras/backend/ones_like>`_
+    """
+    return K.ones_like(*args, **kwargs)
 
 
-def c_to_py_fun(op_name, name = "dataset"):
+def many_replication(grid, replications, axis=0, **kwargs):
     """
-    Map the NNPDF operations to Keras layers
-    NNPDF operations are defined in :py:func:`validphys.convolution.OP
+    Generates a tensor with one extra dimension:
+        a repetition of "grid" n times along the given axis
+    from keras documentation:
+    If x has shape (s1, s2, s3) and axis is 1, the output will have shape (s1, s2 * rep, s3)
+    see full `docs <https://www.tensorflow.org/api_docs/python/tf/keras/backend/repeat_elements>`_
+    """
+    return K.repeat_elements(grid, rep=replications, axis=axis, **kwargs)
 
-    Parameters
-    ----------
-        op_name: str
-            A string defining the operation name
+
+# Property operations
+# modify properties of the tensor like the shape or elements it has
+def flatten(x):
+    """ Flatten tensor x """
+    return tf.reshape(x, (-1,))
+
+
+def boolean_mask(*args, **kwargs):
     """
-    try:
-        operation = OP[op_name]
-    except KeyError as e:
-        raise ValueError(f"Operation {op_name} not recognised") from e
+    Applies a boolean mask to a tensor
 
-    # Convert the operation into a lambda layer
-    operation_layer = keras_Lambda(lambda x: operation(*x), name=f"op_{name}_{op_name}")
-    return operation_layer
+    Relevant parameters: (tensor, mask, axis=None)
+    see full `docs <https://www.tensorflow.org/api_docs/python/tf/boolean_mask>`_.
+    """
+    return tf.boolean_mask(*args, **kwargs)
 
-def op_subtract(o_list): #TODO to be removed, not used once other PRs are merged
-    from tensorflow.keras.layers import subtract
-    return subtract(o_list)
+
+def transpose(tensor, **kwargs):
+    """
+    Transpose a layer,
+    see full `docs <https://www.tensorflow.org/api_docs/python/tf/keras/backend/transpose>`_
+    """
+    return K.transpose(tensor, **kwargs)
+
+
+def concatenate(tensor_list, axis=-1, target_shape=None):
+    """
+    Concatenates a list of numbers or tenosr into a bigger tensor
+    If the target shape is given, the output is reshaped to said shape
+    """
+    concatenated_tensor = K.concatenate(tensor_list, axis=axis)
+    if target_shape:
+        return K.reshape(concatenated_tensor, target_shape)
+    else:
+        return concatenated_tensor
+
+
+# Mathematical operations
+def tensor_product(*args, **kwargs):
+    """
+    Computes the tensordot product between tensor_x and tensor_y
+    See full `docs <https://www.tensorflow.org/api_docs/python/tf/tensordot>`_
+    """
+    return tf.tensordot(*args, **kwargs)
 
 
 def op_multiply(o_list, **kwargs):
@@ -134,3 +217,11 @@ def op_log(o_tensor, **kwargs):
     Computes the logarithm of the input
     """
     return K.log(o_tensor)
+
+
+def sum(*args, **kwargs):
+    """
+    Computes the sum of the elements of the tensor
+    see full `docs <https://www.tensorflow.org/api_docs/python/tf/keras/backend/sum>`_
+    """
+    return K.sum(*args, **kwargs)

n3fit/src/n3fit/layers/DIS.py

@@ -1,5 +1,6 @@
 import numpy as np
 from n3fit.layers.Observable import Observable
+from n3fit.backends import operations as op
 
 
 class DIS(Observable):
@@ -36,8 +37,8 @@ def call(self, pdf_in):
                 - `result`: rank 1 tensor (ndata)
         """
         pdf_in = self.digest_pdf(pdf_in)
-        pdf_masked = self.boolean_mask(pdf_in, self.basis, axis=1)
+        pdf_masked = op.boolean_mask(pdf_in, self.basis, axis=1)
 
-        pdfT = self.transpose(pdf_masked)
-        result = self.tensor_product(self.fktable, pdfT, axes=2)
+        pdfT = op.transpose(pdf_masked)
+        result = op.tensor_product(self.fktable, pdfT, axes=2)
         return result