Multi-replica fits via trvl-mask layers + lru caches in validphys

conda-recipe/meta.yaml

@@ -55,7 +55,8 @@ requirements:
         - pineappl >=0.6.2
         - eko >=0.14.1
         - fiatlux
-        - curio >=1.0 # reportengine uses it but it's not in its dependencies 
+        - frozendict  # needed for caching of data loading
+        - curio >=1.0 # reportengine uses it but it's not in its dependencies
 
 test:
     requires:

n3fit/runcards/reproduce_nnpdf40/NNPDF40_nnlo_as_01180_1000.yml

@@ -163,4 +163,4 @@ integrability:
 
 ############################################################
 debug: false
-maxcores: 4
+maxcores: 4

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

@@ -211,6 +211,12 @@ def flatten(x):
     return tf.reshape(x, (-1,))
 
 
+@tf.function
+def reshape(x, shape):
+    """reshape tensor x"""
+    return tf.reshape(x, shape)
+
+
 def boolean_mask(*args, **kwargs):
     """
     Applies a boolean mask to a tensor

n3fit/src/n3fit/checks.py

@@ -380,11 +380,6 @@ def check_consistent_parallel(parameters, parallel_models, same_trvl_per_replica
     """
     if not parallel_models:
         return
-    if not same_trvl_per_replica:
-        raise CheckError(
-            "Replicas cannot be run in parallel with different training/validation "
-            " masks, please set `same_trvl_per_replica` to True in the runcard"
-        )
     if parameters.get("layer_type") != "dense":
         raise CheckError("Parallelization has only been tested with layer_type=='dense'")
 

n3fit/src/n3fit/layers/losses.py

@@ -8,6 +8,7 @@
 
 """
 import numpy as np
+
 from n3fit.backends import MetaLayer
 from n3fit.backends import operations as op
 
@@ -37,12 +38,11 @@ class LossInvcovmat(MetaLayer):
     True
     """
 
-    def __init__(self, invcovmat, y_true, mask=None, covmat=None, **kwargs):
+    def __init__(self, invcovmat, y_true, mask=None, covmat=None, diag=False, **kwargs):
         # If we have a diagonal matrix, padd with 0s and hope it's not too heavy on memory
-        if len(invcovmat.shape) == 1:
-            invcovmat = np.diag(invcovmat)
         self._invcovmat = op.numpy_to_tensor(invcovmat)
         self._covmat = covmat
+        self._diag = diag
         self._y_true = op.numpy_to_tensor(y_true)
         self._ndata = y_true.shape[-1]
         if mask is None or all(mask):
@@ -56,9 +56,7 @@ def build(self, input_shape):
         """Transform the inverse covmat and the mask into
         weights of the layers"""
         init = MetaLayer.init_constant(self._invcovmat)
-        self.kernel = self.builder_helper(
-            "invcovmat", (self._ndata, self._ndata), init, trainable=False
-        )
+        self.kernel = self.builder_helper("invcovmat", self._invcovmat.shape, init, trainable=False)
         mask_shape = (1, 1, self._ndata)
         if self._mask is None:
             init_mask = MetaLayer.init_constant(np.ones(mask_shape))
@@ -71,7 +69,10 @@ def add_covmat(self, covmat):
         Note, however, that the _covmat attribute of the layer will
         still refer to the original data covmat
         """
-        new_covmat = np.linalg.inv(self._covmat + covmat)
+        if self._diag:
+            new_covmat = np.invert(self._covmat + covmat)
+        else:
+            new_covmat = np.linalg.inv(self._covmat + covmat)
         self.kernel.assign(new_covmat)
 
     def update_mask(self, new_mask):
@@ -82,13 +83,45 @@ def call(self, y_pred, **kwargs):
         tmp_raw = self._y_true - y_pred
         # TODO: most of the time this is a y * I multiplication and can be skipped
         # benchmark how much time (if any) is lost in this in actual fits for the benefit of faster kfolds
+        #        import pdb; pdb.set_trace()
         tmp = op.op_multiply([tmp_raw, self.mask])
+        if self._diag:
+            return LossInvcovmat.contract_covmat_diag(self.kernel, tmp)
+        else:
+            return LossInvcovmat.contract_covmat(self.kernel, tmp)
+
+    @staticmethod
+    def contract_covmat(kernel, tmp):
+        if tmp.shape[1] == 1:
+            # einsum is not well suited for CPU, so use tensordot if not multimodel
+            if len(kernel.shape) == 3:
+                right_dot = op.tensor_product(kernel[0, ...], tmp[0, 0, :], axes=1)
+                res = op.tensor_product(tmp[0, :, :], right_dot, axes=1)
+            else:
+                right_dot = op.tensor_product(kernel, tmp[0, 0, :], axes=1)
+                res = op.tensor_product(tmp[0, :, :], right_dot, axes=1)
+        else:
+            if len(kernel.shape) == 3:
+                res = op.einsum("bri, rij, brj -> r", tmp, kernel, tmp)
+            else:
+                res = op.einsum("bri, ij, brj -> r", tmp, kernel, tmp)
+        return res
+
+    @staticmethod
+    def contract_covmat_diag(kernel, tmp):
         if tmp.shape[1] == 1:
             # einsum is not well suited for CPU, so use tensordot if not multimodel
-            right_dot = op.tensor_product(self.kernel, tmp[0, 0, :], axes=1)
-            res = op.tensor_product(tmp[0, :, :], right_dot, axes=1)
+            if len(kernel.shape) == 2:
+                right_dot = op.tensor_product(kernel[0, :], tmp[0, 0, :], axes=1)
+                res = op.tensor_product(tmp[0, :, :], right_dot, axes=1)
+            else:
+                right_dot = op.tensor_product(kernel, tmp[0, 0, :], axes=1)
+                res = op.tensor_product(tmp[0, :, :], right_dot, axes=1)
         else:
-            res = op.einsum("bri, ij, brj -> r", tmp, self.kernel, tmp)
+            if len(kernel.shape) == 2:
+                res = op.einsum("bri, ri, bri -> r", tmp, kernel, tmp)
+            else:
+                res = op.einsum("bri, i, bri -> r", tmp, kernel, tmp)
         return res
 
 

n3fit/src/n3fit/layers/mask.py

@@ -1,3 +1,5 @@
+from numpy import count_nonzero
+
 from n3fit.backends import MetaLayer
 from n3fit.backends import operations as op
 
@@ -19,29 +21,36 @@ class Mask(MetaLayer):
         c: float
             constant multiplier for every output
         axis: int
-            axis in which to apply the mask
+            axis in which to apply the mask. Currently,
+            only the last axis gives the correct output shape
     """
 
     def __init__(self, bool_mask=None, c=None, axis=None, **kwargs):
         if bool_mask is None:
             self.mask = None
+            self.last_dim = -1
         else:
             self.mask = op.numpy_to_tensor(bool_mask, dtype=bool)
+            if len(bool_mask.shape) == 1:
+                self.last_dim = count_nonzero(bool_mask)
+            else:
+                self.last_dim = count_nonzero(bool_mask[0, ...])
         self.c = c
         self.axis = axis
         super().__init__(**kwargs)
 
     def build(self, input_shape):
         if self.c is not None:
             initializer = MetaLayer.init_constant(value=self.c)
-            self.kernel = self.builder_helper(
-                "mask", (1,), initializer, trainable=False
-            )
+            self.kernel = self.builder_helper("mask", (1,), initializer, trainable=False)
         super(Mask, self).build(input_shape)
 
     def call(self, ret):
         if self.mask is not None:
-            ret = op.boolean_mask(ret, self.mask, axis=self.axis)
+            flat_res = op.boolean_mask(ret, self.mask, axis=self.axis)
+            output_shape = [-1 if d is None else d for d in ret.get_shape()]
+            output_shape[-1] = self.last_dim
+            ret = op.reshape(flat_res, shape=output_shape)
         if self.c is not None:
             ret = ret * self.kernel
         return ret