Polarized proton-PDF fits

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

@@ -308,6 +308,21 @@ def pow(tensor, power):
     return tf.pow(tensor, power)
 
 
+@tf.function
+def absolute(tensor):
+    """
+    Compute the absolute value of a tensor
+    """
+    return K.abs(tensor)
+
+
+def multiply_minusone(tensor):
+    """
+    Multiply the elements of a given tensor by (-1)
+    """
+    return keras_Lambda(lambda x: -1 * x)(tensor)
+
+
 @tf.function(experimental_relax_shapes=True)
 def op_log(o_tensor, **kwargs):
     """

n3fit/src/n3fit/checks.py

@@ -5,12 +5,9 @@
 import numbers
 import os
 
-import numpy as np
-
 from n3fit.hyper_optimization import penalties as penalties_module
 from n3fit.hyper_optimization import rewards as rewards_module
 from reportengine.checks import CheckError, make_argcheck
-from validphys.core import PDF
 from validphys.pdfbases import check_basis
 
 log = logging.getLogger(__name__)
@@ -71,6 +68,18 @@ def check_stopping(parameters):
         raise CheckError(f"Needs to run at least 1 epoch, got: {epochs}")
 
 
+@make_argcheck
+def check_polarised(fitbasis, unpolpdf, fitting):
+    """Checks that if the polarised basis is used, then then
+    the necessary entries are specified correctly.
+    """
+    if "POL" in fitbasis:
+        if unpolpdf is None:
+            raise CheckError("'unpolpdf' needs to be specified for polarised fits.")
+        if fitting.get("sum_rules") != "TSR":
+            raise CheckError("'sum_rules' needs to be 'TSR' for polarised fits.")
+
+
 def check_basis_with_layers(basis, parameters):
     """Check that the last layer matches the number of flavours defined in the runcard"""
     number_of_flavours = len(basis)
@@ -316,7 +325,7 @@ def check_sumrules(sum_rules):
     """Checks that the chosen option for the sum rules are sensible"""
     if isinstance(sum_rules, bool):
         return
-    accepted_options = ["ALL", "MSR", "VSR"]
+    accepted_options = ["ALL", "MSR", "VSR", "TSR"]
     if sum_rules.upper() in accepted_options:
         return
     raise CheckError(f"The only accepted options for the sum rules are: {accepted_options}")

n3fit/src/n3fit/layers/DIS.py

@@ -13,23 +13,23 @@
 
 class DIS(Observable):
     """
-        The DIS class receives a list of active flavours and a fktable
-        and prepares a layer that performs the convolution of said fktable with
-        the incoming pdf.
+    The DIS class receives a list of active flavours and a fktable
+    and prepares a layer that performs the convolution of said fktable with
+    the incoming pdf.
 
-        The fktable is expected to be rank 3 (ndata, xgrid, flavours)
-        while the input pdf is rank 4 where the first dimension is the batch dimension
-        and the last dimension the number of replicas being fitted (1, xgrid, flavours, replicas)
+    The fktable is expected to be rank 3 (ndata, xgrid, flavours)
+    while the input pdf is rank 4 where the first dimension is the batch dimension
+    and the last dimension the number of replicas being fitted (1, xgrid, flavours, replicas)
     """
 
     def gen_mask(self, basis):
         """
-            Receives a list of active flavours and generates a boolean mask tensor
+        Receives a list of active flavours and generates a boolean mask tensor
 
-            Parameters
-            ----------
-                basis: list(int)
-                    list of active flavours
+        Parameters
+        ----------
+            basis: list(int)
+                list of active flavours
         """
         if basis is None:
             self.basis = np.ones(self.nfl, dtype=bool)
@@ -41,37 +41,40 @@ def gen_mask(self, basis):
 
     def call(self, pdf):
         """
-            This function perform the fktable \otimes pdf convolution.
+        This function perform the fktable \otimes pdf convolution.
 
-            First pass the input PDF through a mask to remove the unactive flavors,
-            then a tensor_product between the PDF and each fktable is performed
-            finally the defined operation is applied to all the results
+        First pass the input PDF through a mask to remove the unactive flavors,
+        then a tensor_product between the PDF and each fktable is performed
+        finally the defined operation is applied to all the results
 
-            Parameters
-            ----------
-                pdf:  backend tensor
-                    rank 4 tensor (batch_size, xgrid, flavours, replicas)
+        Parameters
+        ----------
+            pdf:  backend tensor
+                rank 4 tensor (batch_size, xgrid, flavours, replicas)
 
-            Returns
-            -------
-                result: backend tensor
-                    rank 3 tensor (batchsize, replicas, ndata)
+        Returns
+        -------
+            result: backend tensor
+                rank 3 tensor (batchsize, replicas, ndata)
         """
         # DIS never needs splitting
         if self.splitting is not None:
             raise ValueError("DIS layer call with a dataset that needs more than one xgrid?")
 
         results = []
-        # Separate the two possible paths this layer can take
-        if self.many_masks:
-            for mask, fktable in zip(self.all_masks, self.fktables):
+        for idx, fktable in enumerate(self.fktables):
+            mask = self.all_masks[idx] if self.many_masks else self.all_masks[0]
+            if self.check_pol_positivity() and idx == 1:  # Unpolarised POS dataset
+                pdf_masked = op.boolean_mask(self.computed_pdfs[idx], mask, axis=2)
+            else:
                 pdf_masked = op.boolean_mask(pdf, mask, axis=2)
+
+            if self.check_pol_positivity() and idx == 0:  # Polarised POS dataset
                 res = op.tensor_product(pdf_masked, fktable, axes=[(1, 2), (2, 1)])
-                results.append(res)
-        else:
-            pdf_masked = op.boolean_mask(pdf, self.all_masks[0], axis=2)
-            for fktable in self.fktables:
+                res = op.multiply_minusone(op.absolute(res))
+            else:
                 res = op.tensor_product(pdf_masked, fktable, axes=[(1, 2), (2, 1)])
-                results.append(res)
+
+            results.append(res)
 
         return self.operation(results)

n3fit/src/n3fit/layers/msr_normalization.py

@@ -21,28 +21,25 @@
     'v15': 6,
     'v24': 7,
     'v35': 8,
+    't3': 9,
+    't8': 10,
+    't15': 11,
+    't24': 12,
+    't35': 13,
 }
 MSR_COMPONENTS = ['g']
 MSR_DENOMINATORS = {'g': 'g'}
 # The VSR normalization factor of component f is given by
 # VSR_CONSTANTS[f] / VSR_DENOMINATORS[f]
 VSR_COMPONENTS = ['v', 'v35', 'v24', 'v3', 'v8', 'v15']
-VSR_CONSTANTS = {
-    'v': 3.0,
-    'v35': 3.0,
-    'v24': 3.0,
-    'v3': 1.0,
-    'v8': 3.0,
-    'v15': 3.0,
-}
-VSR_DENOMINATORS = {
-    'v': 'v',
-    'v35': 'v',
-    'v24': 'v',
-    'v3': 'v3',
-    'v8': 'v8',
-    'v15': 'v15',
-}
+VSR_CONSTANTS = {'v': 3.0, 'v35': 3.0, 'v24': 3.0, 'v3': 1.0, 'v8': 3.0, 'v15': 3.0}
+VSR_DENOMINATORS = {'v': 'v', 'v35': 'v', 'v24': 'v', 'v3': 'v3', 'v8': 'v8', 'v15': 'v15'}
+
+# The following lays out the SR for Polarised PDFs
+TSR_COMPONENTS = ['t3', 't8']
+TSR_DENOMINATORS = {'t3': 't3', 't8': 't8'}
+# TODO: Find better way for account for the uncertainties
+TSR_CONSTANTS = {'t3': 1.2701, 't8': 0.5850}
 
 
 class MSR_Normalization(MetaLayer):
@@ -52,6 +49,7 @@ class MSR_Normalization(MetaLayer):
 
     _msr_enabled = False
     _vsr_enabled = False
+    _tsr_enabled = False
 
     def __init__(self, mode: str = "ALL", replicas: int = 1, **kwargs):
         if mode == True or mode.upper() == "ALL":
@@ -61,6 +59,8 @@ def __init__(self, mode: str = "ALL", replicas: int = 1, **kwargs):
             self._msr_enabled = True
         elif mode.upper() == "VSR":
             self._vsr_enabled = True
+        elif mode.upper() == "TSR":
+            self._tsr_enabled = True
         else:
             raise ValueError(f"Mode {mode} not accepted for sum rules")
 
@@ -75,6 +75,13 @@ def __init__(self, mode: str = "ALL", replicas: int = 1, **kwargs):
             self.vsr_factors = op.numpy_to_tensor(
                 [np.repeat(VSR_CONSTANTS[c], replicas) for c in VSR_COMPONENTS]
             )
+        if self._tsr_enabled:
+            self.divisor_indices += [IDX[TSR_DENOMINATORS[c]] for c in TSR_COMPONENTS]
+            indices += [IDX[c] for c in TSR_COMPONENTS]
+            self.tsr_factors = op.numpy_to_tensor(
+                [np.repeat(TSR_CONSTANTS[c], replicas) for c in TSR_COMPONENTS]
+            )
+
         # Need this extra dimension for the scatter_to_one operation
         self.indices = [[i] for i in indices]
 
@@ -113,6 +120,8 @@ def call(self, pdf_integrated, photon_integral):
             ]
         if self._vsr_enabled:
             numerators += [self.vsr_factors]
+        if self._tsr_enabled:
+            numerators += [self.tsr_factors]
 
         numerators = op.concatenate(numerators, axis=0)
         divisors = op.gather(y, self.divisor_indices, axis=0)

n3fit/src/n3fit/layers/observable.py

@@ -38,11 +38,25 @@ class Observable(MetaLayer, ABC):
             number of flavours in the pdf (default:14)
     """
 
-    def __init__(self, fktable_data, fktable_arr, operation_name, nfl=14, **kwargs):
+    def __init__(
+        self,
+        fktable_data,
+        fktable_arr,
+        dataset_name,
+        fitbasis,
+        extern_lhapdf,
+        operation_name,
+        nfl=14,
+        **kwargs
+    ):
         super(MetaLayer, self).__init__(**kwargs)
 
+        self.dataset_name = dataset_name
         self.nfl = nfl
+        self.fitbasis = fitbasis
+        self.fktable_data = fktable_data
 
+        self.computed_pdfs = []
         basis = []
         xgrids = []
         self.fktables = []
@@ -51,6 +65,13 @@ def __init__(self, fktable_data, fktable_arr, operation_name, nfl=14, **kwargs):
             basis.append(fkdata.luminosity_mapping)
             self.fktables.append(op.numpy_to_tensor(fk))
 
+            if self.check_pol_positivity():
+                resx = extern_lhapdf(fkdata.xgrid.tolist())
+                resx = np.expand_dims(resx, axis=[0, -1])
+                self.computed_pdfs.append(op.numpy_to_tensor(resx))
+                # TODO: Apply the following systematically from `vp`
+                operation_name = "ADD"
+
         # check how many xgrids this dataset needs
         if is_unique(xgrids):
             self.splitting = None
@@ -71,6 +92,12 @@ def __init__(self, fktable_data, fktable_arr, operation_name, nfl=14, **kwargs):
     def compute_output_shape(self, input_shape):
         return (self.output_dim, None)
 
+    def check_pol_positivity(self):
+        if "POL" in self.fitbasis and "_POS_" in self.dataset_name:
+            if len(self.fktable_data) == 2:
+                return True
+        return False
+
     # Overridables
     @abstractmethod
     def gen_mask(self, basis):

n3fit/src/n3fit/model_gen.py

@@ -99,7 +99,7 @@ def __call__(self, pdf_layer, mask=None):
 
 
 def observable_generator(
-    spec_dict, positivity_initial=1.0, integrability=False
+    spec_dict, fitbasis, extern_lhapdf, positivity_initial=1.0, integrability=False
 ):  # pylint: disable=too-many-locals
     """
     This function generates the observable models for each experiment.
@@ -178,6 +178,9 @@ def observable_generator(
             obs_layer_tr = Obs_Layer(
                 dataset.fktables_data,
                 dataset.training_fktables(),
+                dataset_name,
+                fitbasis,
+                extern_lhapdf,
                 operation_name,
                 name=f"dat_{dataset_name}",
             )
@@ -187,6 +190,9 @@ def observable_generator(
             obs_layer_ex = Obs_Layer(
                 dataset.fktables_data,
                 dataset.fktables(),
+                dataset_name,
+                fitbasis,
+                extern_lhapdf,
                 operation_name,
                 name=f"exp_{dataset_name}",
             )
@@ -195,18 +201,27 @@ def observable_generator(
             obs_layer_tr = Obs_Layer(
                 dataset.fktables_data,
                 dataset.training_fktables(),
+                dataset_name,
+                fitbasis,
+                extern_lhapdf,
                 operation_name,
                 name=f"dat_{dataset_name}",
             )
             obs_layer_ex = Obs_Layer(
                 dataset.fktables_data,
                 dataset.fktables(),
+                dataset_name,
+                fitbasis,
+                extern_lhapdf,
                 operation_name,
                 name=f"exp_{dataset_name}",
             )
             obs_layer_vl = Obs_Layer(
                 dataset.fktables_data,
                 dataset.validation_fktables(),
+                dataset_name,
+                fitbasis,
+                extern_lhapdf,
                 operation_name,
                 name=f"val_{dataset_name}",
             )