From cc57f4ac7dcb1817072dc14327d323831465b5a4 Mon Sep 17 00:00:00 2001 From: Aron Date: Thu, 13 Jul 2023 14:51:14 +0200 Subject: [PATCH 01/21] Add test of MSR_Normalization layer --- n3fit/src/n3fit/tests/test_msr.py | 84 +++++++++++++++++++++++++++++++ 1 file changed, 84 insertions(+) create mode 100644 n3fit/src/n3fit/tests/test_msr.py diff --git a/n3fit/src/n3fit/tests/test_msr.py b/n3fit/src/n3fit/tests/test_msr.py new file mode 100644 index 0000000000..993aabe940 --- /dev/null +++ b/n3fit/src/n3fit/tests/test_msr.py @@ -0,0 +1,84 @@ +import numpy as np + +from n3fit.backends import operations as op +from n3fit.layers import MSR_Normalization + + +def test_layer(layer): + np.random.seed(422) + pdf_integrated = op.numpy_to_tensor(np.random.normal(size=(1, 14))) + photon_integral = op.numpy_to_tensor(np.random.normal(size=(1,))) + + return layer(pdf_integrated, photon_integral) + + +def test_all(): + layer = MSR_Normalization(mode='ALL') + output = test_layer(layer) + known_output = op.numpy_to_tensor( + [ + 1.0, + 1.0, + 6.7740397, + -3.8735497, + 15.276561, + 5.9522753, + 8.247783, + -3.8735497, + -3.8735497, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + ] + ) + np.testing.assert_allclose(output, known_output, rtol=1e-5) + + +def test_msr(): + layer = MSR_Normalization(mode='MSR') + output = test_layer(layer) + known_output = op.numpy_to_tensor( + [ + 1.0, + 1.0, + 6.7740397, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + ] + ) + np.testing.assert_allclose(output, known_output, rtol=1e-5) + + +def test_vsr(): + layer = MSR_Normalization(mode='VSR') + output = test_layer(layer) + known_output = op.numpy_to_tensor( + [ + 1.0, + 1.0, + 1.0, + -3.8735497, + 15.276561, + 5.9522753, + 8.247783, + -3.8735497, + -3.8735497, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + ] + ) + np.testing.assert_allclose(output, known_output, rtol=1e-5) From fa5333f44c6e8c04281f31564a486ec5ffad627d Mon Sep 17 00:00:00 2001 From: Aron Date: Tue, 18 Jul 2023 10:39:00 +0200 Subject: [PATCH 02/21] Add photon integral before MSR_Normalization layer --- n3fit/src/n3fit/layers/msr_normalization.py | 6 ++---- n3fit/src/n3fit/msr.py | 12 ++++++++---- n3fit/src/n3fit/tests/test_msr.py | 9 +++++++-- 3 files changed, 17 insertions(+), 10 deletions(-) diff --git a/n3fit/src/n3fit/layers/msr_normalization.py b/n3fit/src/n3fit/layers/msr_normalization.py index 552f618b59..6e8bc98875 100644 --- a/n3fit/src/n3fit/layers/msr_normalization.py +++ b/n3fit/src/n3fit/layers/msr_normalization.py @@ -39,7 +39,7 @@ def __init__(self, output_dim=14, mode="ALL", **kwargs): super().__init__(**kwargs) - def call(self, pdf_integrated, photon_integral): + def call(self, pdf_integrated): """Imposes the valence and momentum sum rules: A_g = (1-sigma-photon)/g A_v = A_v24 = A_v35 = 3/V @@ -53,8 +53,6 @@ def call(self, pdf_integrated, photon_integral): ---------- pdf_integrated: (Tensor(1,None,14)) the integrated PDF - photon_integral: (Tensor(1)): - the integrated photon, not included in PDF Returns ------- @@ -65,7 +63,7 @@ def call(self, pdf_integrated, photon_integral): norm_constants = [] if self._msr_enabled: - n_ag = [(1.0 - y[GLUON_IDX[0][0] - 1] - photon_integral[0]) / y[GLUON_IDX[0][0]]] * len( + n_ag = [(1.0 - y[GLUON_IDX[0][0] - 1] - y[GLUON_IDX[0][0] - 2]) / y[GLUON_IDX[0][0]]] * len( GLUON_IDX ) norm_constants += n_ag diff --git a/n3fit/src/n3fit/msr.py b/n3fit/src/n3fit/msr.py index 17392d3022..532d0421bd 100644 --- a/n3fit/src/n3fit/msr.py +++ b/n3fit/src/n3fit/msr.py @@ -86,13 +86,17 @@ def generate_msr_model_and_grid( pdf_integrated = xIntegrator(weights_array, input_shape=(nx,))(pdf_integrand) # 5. THe input for the photon integral, will be set to 0 if no photons - photon_integral = Input(shape=(), batch_size=1, name='photon_integral') + photon_integral = Input(shape=(1,), batch_size=1, name='photon_integral') + + # 5b. Insert the photon integral as the first component of the pdf integrals + pdf_integrated = Lambda( + lambda pdf_photon: op.stack([pdf_photon[1], pdf_photon[0][:, 1:]], axis=1), + name="join_photon", + )([pdf_integrated, photon_integral]) # 6. Compute the normalization factor # For now set the photon component to None - normalization_factor = MSR_Normalization(output_dim, mode, name="msr_weights")( - pdf_integrated, photon_integral - ) + normalization_factor = MSR_Normalization(output_dim, mode, name="msr_weights")(pdf_integrated) # 7. Apply the normalization factor to the pdf pdf_normalized = Lambda(lambda pdf_norm: pdf_norm[0] * pdf_norm[1], name="pdf_normalized")( diff --git a/n3fit/src/n3fit/tests/test_msr.py b/n3fit/src/n3fit/tests/test_msr.py index 993aabe940..77b86aab3e 100644 --- a/n3fit/src/n3fit/tests/test_msr.py +++ b/n3fit/src/n3fit/tests/test_msr.py @@ -7,9 +7,14 @@ def test_layer(layer): np.random.seed(422) pdf_integrated = op.numpy_to_tensor(np.random.normal(size=(1, 14))) - photon_integral = op.numpy_to_tensor(np.random.normal(size=(1,))) - return layer(pdf_integrated, photon_integral) + # done this way to stay consistent with the original code + photon_integral = op.numpy_to_tensor([np.random.normal(size=(1,))]) + pdf_integrated = op.concatenate( + [op.numpy_to_tensor(photon_integral), pdf_integrated[:, 1:]], axis=1 + ) + + return layer(pdf_integrated) def test_all(): From ab4bb632356b34c2708e8f3b94516e58b58d9a79 Mon Sep 17 00:00:00 2001 From: Aron Date: Tue, 18 Jul 2023 10:50:58 +0200 Subject: [PATCH 03/21] Simplify indices --- n3fit/src/n3fit/layers/msr_normalization.py | 34 +++++++++++++-------- 1 file changed, 21 insertions(+), 13 deletions(-) diff --git a/n3fit/src/n3fit/layers/msr_normalization.py b/n3fit/src/n3fit/layers/msr_normalization.py index 6e8bc98875..daa607d50c 100644 --- a/n3fit/src/n3fit/layers/msr_normalization.py +++ b/n3fit/src/n3fit/layers/msr_normalization.py @@ -1,11 +1,20 @@ from n3fit.backends import MetaLayer from n3fit.backends import operations as op -GLUON_IDX = [[2]] V_IDX = [[3], [7], [8]] -V3_IDX = [[4]] -V8_IDX = [[5]] -V15_IDX = [[6]] + +indices = { + 'photon': 0, + 'sigma': 1, + 'g': 2, + 'v': 3, + 'v3': 4, + 'v8': 5, + 'v15': 6, + 'v35': 7, + 'v24': 8, +} + class MSR_Normalization(MetaLayer): @@ -29,9 +38,10 @@ def __init__(self, output_dim=14, mode="ALL", **kwargs): idx = [] if self._msr_enabled: - idx += GLUON_IDX + idx += [indices['g']] if self._vsr_enabled: - idx += V_IDX + V3_IDX + V8_IDX + V15_IDX + idx += [indices[f] for f in ['v', 'v35', 'v24', 'v3', 'v8', 'v15']] + idx = [[i] for i in idx] self._out_scatter = op.as_layer( op.scatter_to_one, op_kwargs={"indices": idx, "output_dim": output_dim} @@ -63,16 +73,14 @@ def call(self, pdf_integrated): norm_constants = [] if self._msr_enabled: - n_ag = [(1.0 - y[GLUON_IDX[0][0] - 1] - y[GLUON_IDX[0][0] - 2]) / y[GLUON_IDX[0][0]]] * len( - GLUON_IDX - ) + n_ag = [(1.0 - y[indices['sigma']] - y[indices['photon']]) / y[indices['g']]] norm_constants += n_ag if self._vsr_enabled: - n_av = [3.0 / y[V_IDX[0][0]]] * len(V_IDX) - n_av3 = [1.0 / y[V3_IDX[0][0]]] * len(V3_IDX) - n_av8 = [3.0 / y[V8_IDX[0][0]]] * len(V8_IDX) - n_av15 = [3.0 / y[V15_IDX[0][0]]] * len(V15_IDX) + n_av = [3.0 / y[indices['v']]] * 3 + n_av3 = [1.0 / y[indices['v3']]] + n_av8 = [3.0 / y[indices['v8']]] + n_av15 = [3.0 / y[indices['v15']]] norm_constants += n_av + n_av3 + n_av8 + n_av15 return self._out_scatter(norm_constants) From 83c1f6533130d96dfccfcb5f4cb7dd260538dd90 Mon Sep 17 00:00:00 2001 From: Aron Date: Tue, 18 Jul 2023 10:52:02 +0200 Subject: [PATCH 04/21] Correct shape in docstring --- n3fit/src/n3fit/layers/msr_normalization.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/n3fit/src/n3fit/layers/msr_normalization.py b/n3fit/src/n3fit/layers/msr_normalization.py index daa607d50c..4c88b633dd 100644 --- a/n3fit/src/n3fit/layers/msr_normalization.py +++ b/n3fit/src/n3fit/layers/msr_normalization.py @@ -1,7 +1,6 @@ from n3fit.backends import MetaLayer from n3fit.backends import operations as op -V_IDX = [[3], [7], [8]] indices = { 'photon': 0, @@ -16,7 +15,6 @@ } - class MSR_Normalization(MetaLayer): """ Applies the normalisation so that the PDF output fullfills the sum rules @@ -61,7 +59,7 @@ def call(self, pdf_integrated): Parameters ---------- - pdf_integrated: (Tensor(1,None,14)) + pdf_integrated: (Tensor(1, 14)) the integrated PDF Returns From d1482c824576ae9760f1326e7aee2d5026f7da31 Mon Sep 17 00:00:00 2001 From: Aron Date: Tue, 18 Jul 2023 11:37:27 +0200 Subject: [PATCH 05/21] Change notation for indices --- n3fit/src/n3fit/layers/msr_normalization.py | 22 ++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/n3fit/src/n3fit/layers/msr_normalization.py b/n3fit/src/n3fit/layers/msr_normalization.py index 4c88b633dd..baaebf1753 100644 --- a/n3fit/src/n3fit/layers/msr_normalization.py +++ b/n3fit/src/n3fit/layers/msr_normalization.py @@ -2,7 +2,7 @@ from n3fit.backends import operations as op -indices = { +IDX = { 'photon': 0, 'sigma': 1, 'g': 2, @@ -34,15 +34,15 @@ def __init__(self, output_dim=14, mode="ALL", **kwargs): else: raise ValueError(f"Mode {mode} not accepted for sum rules") - idx = [] + indices = [] if self._msr_enabled: - idx += [indices['g']] + indices += [IDX['g']] if self._vsr_enabled: - idx += [indices[f] for f in ['v', 'v35', 'v24', 'v3', 'v8', 'v15']] - idx = [[i] for i in idx] + indices += [IDX[f] for f in ['v', 'v35', 'v24', 'v3', 'v8', 'v15']] + indices = [[i] for i in indices] self._out_scatter = op.as_layer( - op.scatter_to_one, op_kwargs={"indices": idx, "output_dim": output_dim} + op.scatter_to_one, op_kwargs={"indices": indices, "output_dim": output_dim} ) super().__init__(**kwargs) @@ -71,14 +71,14 @@ def call(self, pdf_integrated): norm_constants = [] if self._msr_enabled: - n_ag = [(1.0 - y[indices['sigma']] - y[indices['photon']]) / y[indices['g']]] + n_ag = [(1.0 - y[IDX['sigma']] - y[IDX['photon']]) / y[IDX['g']]] norm_constants += n_ag if self._vsr_enabled: - n_av = [3.0 / y[indices['v']]] * 3 - n_av3 = [1.0 / y[indices['v3']]] - n_av8 = [3.0 / y[indices['v8']]] - n_av15 = [3.0 / y[indices['v15']]] + n_av = [3.0 / y[IDX['v']]] * 3 + n_av3 = [1.0 / y[IDX['v3']]] + n_av8 = [3.0 / y[IDX['v8']]] + n_av15 = [3.0 / y[IDX['v15']]] norm_constants += n_av + n_av3 + n_av8 + n_av15 return self._out_scatter(norm_constants) From 073df14118adfdb02bd5ab684827c1113ac9bed6 Mon Sep 17 00:00:00 2001 From: Aron Date: Tue, 18 Jul 2023 11:55:29 +0200 Subject: [PATCH 06/21] Replace flatten with [0] to get rid of batch dimension --- n3fit/src/n3fit/layers/msr_normalization.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/n3fit/src/n3fit/layers/msr_normalization.py b/n3fit/src/n3fit/layers/msr_normalization.py index baaebf1753..7315d1a86c 100644 --- a/n3fit/src/n3fit/layers/msr_normalization.py +++ b/n3fit/src/n3fit/layers/msr_normalization.py @@ -67,7 +67,7 @@ def call(self, pdf_integrated): normalization_factor: Tensor(14) The normalization factors per flavour. """ - y = op.flatten(pdf_integrated) + y = pdf_integrated[0] # get rid of the batch dimension norm_constants = [] if self._msr_enabled: From 36c2fff9b20939d40f35ca523362939f5012bc4c Mon Sep 17 00:00:00 2001 From: Aron Date: Tue, 18 Jul 2023 12:02:24 +0200 Subject: [PATCH 07/21] Change output_dim to automatically inferred out_shape --- .../backends/keras_backend/operations.py | 4 ++-- n3fit/src/n3fit/layers/msr_normalization.py | 19 +++++++++++-------- 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/n3fit/src/n3fit/backends/keras_backend/operations.py b/n3fit/src/n3fit/backends/keras_backend/operations.py index efff803b92..43a89ea04b 100644 --- a/n3fit/src/n3fit/backends/keras_backend/operations.py +++ b/n3fit/src/n3fit/backends/keras_backend/operations.py @@ -338,12 +338,12 @@ def split(*args, **kwargs): return tf.split(*args, **kwargs) -def scatter_to_one(values, indices=[[1]], output_dim=14): +def scatter_to_one(values, indices, output_shape): """ Like scatter_nd initialized to one instead of zero see full `docs `_ """ - ones = np.ones(output_dim, dtype=np.float32) + ones = np.ones(output_shape, dtype=np.float32) return tf.tensor_scatter_nd_update(ones, indices, values) diff --git a/n3fit/src/n3fit/layers/msr_normalization.py b/n3fit/src/n3fit/layers/msr_normalization.py index 7315d1a86c..217c8cc858 100644 --- a/n3fit/src/n3fit/layers/msr_normalization.py +++ b/n3fit/src/n3fit/layers/msr_normalization.py @@ -34,19 +34,22 @@ def __init__(self, output_dim=14, mode="ALL", **kwargs): else: raise ValueError(f"Mode {mode} not accepted for sum rules") - indices = [] + self.indices = [] if self._msr_enabled: - indices += [IDX['g']] + self.indices += [IDX['g']] if self._vsr_enabled: - indices += [IDX[f] for f in ['v', 'v35', 'v24', 'v3', 'v8', 'v15']] - indices = [[i] for i in indices] - - self._out_scatter = op.as_layer( - op.scatter_to_one, op_kwargs={"indices": indices, "output_dim": output_dim} - ) + self.indices += [IDX[f] for f in ['v', 'v35', 'v24', 'v3', 'v8', 'v15']] + self.indices = [[i] for i in self.indices] super().__init__(**kwargs) + def build(self, input_shape): + self.out_shape = input_shape[1:] + self._out_scatter = lambda pdf_integrated: op.scatter_to_one( + pdf_integrated, indices=self.indices, output_shape=self.out_shape + ) + super().build(input_shape) + def call(self, pdf_integrated): """Imposes the valence and momentum sum rules: A_g = (1-sigma-photon)/g From 50fb75a2d27b5809544af30b0d94e7e016c7a85c Mon Sep 17 00:00:00 2001 From: Aron Date: Tue, 18 Jul 2023 12:22:02 +0200 Subject: [PATCH 08/21] Fix bug with photon --- n3fit/src/n3fit/model_gen.py | 3 ++- n3fit/src/n3fit/msr.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py index 2bd898021b..34e4ce0ee6 100644 --- a/n3fit/src/n3fit/model_gen.py +++ b/n3fit/src/n3fit/model_gen.py @@ -680,8 +680,9 @@ def compute_unnormalized_pdf(x, neural_network, compute_preprocessing_factor): "pdf_xgrid_integration": pdf_integration_grid, "xgrid_integration": integrator_input, # The photon is treated separately, need to get its integrals to normalize the pdf - "photon_integral": op.numpy_to_tensor( + "photon_integral": op.numpy_to_tensor([[ 0.0 if not photons else photons.integral[i_replica] + ]] ), } ) diff --git a/n3fit/src/n3fit/msr.py b/n3fit/src/n3fit/msr.py index 532d0421bd..b04732c1b1 100644 --- a/n3fit/src/n3fit/msr.py +++ b/n3fit/src/n3fit/msr.py @@ -90,7 +90,7 @@ def generate_msr_model_and_grid( # 5b. Insert the photon integral as the first component of the pdf integrals pdf_integrated = Lambda( - lambda pdf_photon: op.stack([pdf_photon[1], pdf_photon[0][:, 1:]], axis=1), + lambda pdf_photon: op.concatenate([pdf_photon[1], pdf_photon[0][:, 1:]], axis=1), name="join_photon", )([pdf_integrated, photon_integral]) From c1dd92518c953babf135203c8ceef508e66abdf4 Mon Sep 17 00:00:00 2001 From: Aron Date: Tue, 18 Jul 2023 14:23:11 +0200 Subject: [PATCH 09/21] Rename test_layer so its not seen as a test --- n3fit/src/n3fit/tests/test_msr.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/n3fit/src/n3fit/tests/test_msr.py b/n3fit/src/n3fit/tests/test_msr.py index 77b86aab3e..25547ec22f 100644 --- a/n3fit/src/n3fit/tests/test_msr.py +++ b/n3fit/src/n3fit/tests/test_msr.py @@ -4,7 +4,7 @@ from n3fit.layers import MSR_Normalization -def test_layer(layer): +def apply_layer_to_fixed_input(layer): np.random.seed(422) pdf_integrated = op.numpy_to_tensor(np.random.normal(size=(1, 14))) @@ -19,7 +19,7 @@ def test_layer(layer): def test_all(): layer = MSR_Normalization(mode='ALL') - output = test_layer(layer) + output = apply_layer_to_fixed_input(layer) known_output = op.numpy_to_tensor( [ 1.0, @@ -43,7 +43,7 @@ def test_all(): def test_msr(): layer = MSR_Normalization(mode='MSR') - output = test_layer(layer) + output = apply_layer_to_fixed_input(layer) known_output = op.numpy_to_tensor( [ 1.0, @@ -67,7 +67,7 @@ def test_msr(): def test_vsr(): layer = MSR_Normalization(mode='VSR') - output = test_layer(layer) + output = apply_layer_to_fixed_input(layer) known_output = op.numpy_to_tensor( [ 1.0, From a25bb5723cd58a82c61700e8ba037c9d9ca20dc2 Mon Sep 17 00:00:00 2001 From: Aron Date: Thu, 27 Jul 2023 13:25:29 +0200 Subject: [PATCH 10/21] Remove unused output_dim argument and out_shape attribute --- n3fit/src/n3fit/layers/msr_normalization.py | 7 +++---- n3fit/src/n3fit/msr.py | 2 +- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/n3fit/src/n3fit/layers/msr_normalization.py b/n3fit/src/n3fit/layers/msr_normalization.py index 217c8cc858..07ca2085d9 100644 --- a/n3fit/src/n3fit/layers/msr_normalization.py +++ b/n3fit/src/n3fit/layers/msr_normalization.py @@ -1,7 +1,6 @@ from n3fit.backends import MetaLayer from n3fit.backends import operations as op - IDX = { 'photon': 0, 'sigma': 1, @@ -23,7 +22,7 @@ class MSR_Normalization(MetaLayer): _msr_enabled = False _vsr_enabled = False - def __init__(self, output_dim=14, mode="ALL", **kwargs): + def __init__(self, mode="ALL", **kwargs): if mode == True or mode.upper() == "ALL": self._msr_enabled = True self._vsr_enabled = True @@ -44,9 +43,9 @@ def __init__(self, output_dim=14, mode="ALL", **kwargs): super().__init__(**kwargs) def build(self, input_shape): - self.out_shape = input_shape[1:] + out_shape.out_shape = input_shape[1:] self._out_scatter = lambda pdf_integrated: op.scatter_to_one( - pdf_integrated, indices=self.indices, output_shape=self.out_shape + pdf_integrated, indices=self.indices, output_shape=out_shap ) super().build(input_shape) diff --git a/n3fit/src/n3fit/msr.py b/n3fit/src/n3fit/msr.py index b04732c1b1..474ec7854a 100644 --- a/n3fit/src/n3fit/msr.py +++ b/n3fit/src/n3fit/msr.py @@ -96,7 +96,7 @@ def generate_msr_model_and_grid( # 6. Compute the normalization factor # For now set the photon component to None - normalization_factor = MSR_Normalization(output_dim, mode, name="msr_weights")(pdf_integrated) + normalization_factor = MSR_Normalization(mode, name="msr_weights")(pdf_integrated) # 7. Apply the normalization factor to the pdf pdf_normalized = Lambda(lambda pdf_norm: pdf_norm[0] * pdf_norm[1], name="pdf_normalized")( From ac8be5126975ef8cfa576c82c2002728dccdd7d3 Mon Sep 17 00:00:00 2001 From: Aron Date: Thu, 27 Jul 2023 13:28:02 +0200 Subject: [PATCH 11/21] Fix order of v24, v35 --- n3fit/src/n3fit/layers/msr_normalization.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/n3fit/src/n3fit/layers/msr_normalization.py b/n3fit/src/n3fit/layers/msr_normalization.py index 07ca2085d9..8c1046c96d 100644 --- a/n3fit/src/n3fit/layers/msr_normalization.py +++ b/n3fit/src/n3fit/layers/msr_normalization.py @@ -9,8 +9,8 @@ 'v3': 4, 'v8': 5, 'v15': 6, - 'v35': 7, - 'v24': 8, + 'v24': 7, + 'v35': 8, } From db2ac8b5c419836bcfed9cd510baca87eae03e64 Mon Sep 17 00:00:00 2001 From: Aron Date: Thu, 27 Jul 2023 14:39:31 +0200 Subject: [PATCH 12/21] Fix typo on out_shape --- n3fit/src/n3fit/layers/msr_normalization.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/n3fit/src/n3fit/layers/msr_normalization.py b/n3fit/src/n3fit/layers/msr_normalization.py index 8c1046c96d..b568088aca 100644 --- a/n3fit/src/n3fit/layers/msr_normalization.py +++ b/n3fit/src/n3fit/layers/msr_normalization.py @@ -43,9 +43,9 @@ def __init__(self, mode="ALL", **kwargs): super().__init__(**kwargs) def build(self, input_shape): - out_shape.out_shape = input_shape[1:] + out_shape = input_shape[1:] self._out_scatter = lambda pdf_integrated: op.scatter_to_one( - pdf_integrated, indices=self.indices, output_shape=out_shap + pdf_integrated, indices=self.indices, output_shape=out_shape ) super().build(input_shape) From 8dc38143d872abb646166941767af40d151a01d1 Mon Sep 17 00:00:00 2001 From: Aron Date: Mon, 7 Aug 2023 14:32:35 +0200 Subject: [PATCH 13/21] Define MSR and VSR indices as global constants --- n3fit/src/n3fit/layers/msr_normalization.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/n3fit/src/n3fit/layers/msr_normalization.py b/n3fit/src/n3fit/layers/msr_normalization.py index b568088aca..0c8502ab79 100644 --- a/n3fit/src/n3fit/layers/msr_normalization.py +++ b/n3fit/src/n3fit/layers/msr_normalization.py @@ -12,6 +12,8 @@ 'v24': 7, 'v35': 8, } +MSR_INDICES = [IDX['g']] +VSR_INDICES = [IDX[f] for f in ['v', 'v35', 'v24', 'v3', 'v8', 'v15']] class MSR_Normalization(MetaLayer): @@ -35,9 +37,9 @@ def __init__(self, mode="ALL", **kwargs): self.indices = [] if self._msr_enabled: - self.indices += [IDX['g']] + self.indices += MSR_INDICES if self._vsr_enabled: - self.indices += [IDX[f] for f in ['v', 'v35', 'v24', 'v3', 'v8', 'v15']] + self.indices += VSR_INDICES self.indices = [[i] for i in self.indices] super().__init__(**kwargs) From a418e99867a0163cb5c5976aeef34172fa7c1723 Mon Sep 17 00:00:00 2001 From: Aron Date: Mon, 7 Aug 2023 14:45:55 +0200 Subject: [PATCH 14/21] Remove build method --- n3fit/src/n3fit/layers/msr_normalization.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/n3fit/src/n3fit/layers/msr_normalization.py b/n3fit/src/n3fit/layers/msr_normalization.py index 0c8502ab79..b1860855f5 100644 --- a/n3fit/src/n3fit/layers/msr_normalization.py +++ b/n3fit/src/n3fit/layers/msr_normalization.py @@ -44,13 +44,6 @@ def __init__(self, mode="ALL", **kwargs): super().__init__(**kwargs) - def build(self, input_shape): - out_shape = input_shape[1:] - self._out_scatter = lambda pdf_integrated: op.scatter_to_one( - pdf_integrated, indices=self.indices, output_shape=out_shape - ) - super().build(input_shape) - def call(self, pdf_integrated): """Imposes the valence and momentum sum rules: A_g = (1-sigma-photon)/g @@ -85,4 +78,8 @@ def call(self, pdf_integrated): n_av15 = [3.0 / y[IDX['v15']]] norm_constants += n_av + n_av3 + n_av8 + n_av15 - return self._out_scatter(norm_constants) + norm_constants = op.scatter_to_one( + norm_constants, indices=self.indices, output_shape=y.shape + ) + + return norm_constants From 9685827d45f1f056675f1f17bd394da772e4d6bb Mon Sep 17 00:00:00 2001 From: Aron Date: Mon, 7 Aug 2023 14:50:44 +0200 Subject: [PATCH 15/21] Add/improve documentation --- n3fit/src/n3fit/layers/msr_normalization.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/n3fit/src/n3fit/layers/msr_normalization.py b/n3fit/src/n3fit/layers/msr_normalization.py index b1860855f5..b4aa056d56 100644 --- a/n3fit/src/n3fit/layers/msr_normalization.py +++ b/n3fit/src/n3fit/layers/msr_normalization.py @@ -18,7 +18,7 @@ class MSR_Normalization(MetaLayer): """ - Applies the normalisation so that the PDF output fullfills the sum rules + Computes the normalisation factors for the sum rules of the PDFs. """ _msr_enabled = False @@ -40,12 +40,14 @@ def __init__(self, mode="ALL", **kwargs): self.indices += MSR_INDICES if self._vsr_enabled: self.indices += VSR_INDICES + # Need this extra dimension for the scatter_to_one operation self.indices = [[i] for i in self.indices] super().__init__(**kwargs) def call(self, pdf_integrated): - """Imposes the valence and momentum sum rules: + """ + Computes the normalization factors for the PDFs: A_g = (1-sigma-photon)/g A_v = A_v24 = A_v35 = 3/V A_v3 = 1/V_3 @@ -78,6 +80,7 @@ def call(self, pdf_integrated): n_av15 = [3.0 / y[IDX['v15']]] norm_constants += n_av + n_av3 + n_av8 + n_av15 + # Fill in the rest of the flavours with 1 norm_constants = op.scatter_to_one( norm_constants, indices=self.indices, output_shape=y.shape ) From bc41e61bdf24a224c02591bdf2da22c8c85e678d Mon Sep 17 00:00:00 2001 From: Aron Date: Tue, 8 Aug 2023 11:57:35 +0200 Subject: [PATCH 16/21] factor out definitions of normalization to top as global constants --- .../backends/keras_backend/operations.py | 7 +++ n3fit/src/n3fit/layers/msr_normalization.py | 45 ++++++++++++++----- 2 files changed, 41 insertions(+), 11 deletions(-) diff --git a/n3fit/src/n3fit/backends/keras_backend/operations.py b/n3fit/src/n3fit/backends/keras_backend/operations.py index 43a89ea04b..a1172902d5 100644 --- a/n3fit/src/n3fit/backends/keras_backend/operations.py +++ b/n3fit/src/n3fit/backends/keras_backend/operations.py @@ -179,6 +179,13 @@ def tmp(x): return layer_op(tensor) +def gather(*args, **kwargs): + """ + Gather elements from a tensor along an axis + """ + return tf.gather(*args, **kwargs) + + # # Tensor operations # f(x: tensor[s]) -> y: tensor diff --git a/n3fit/src/n3fit/layers/msr_normalization.py b/n3fit/src/n3fit/layers/msr_normalization.py index b4aa056d56..ab2d108511 100644 --- a/n3fit/src/n3fit/layers/msr_normalization.py +++ b/n3fit/src/n3fit/layers/msr_normalization.py @@ -12,8 +12,29 @@ 'v24': 7, 'v35': 8, } -MSR_INDICES = [IDX['g']] -VSR_INDICES = [IDX[f] for f in ['v', 'v35', 'v24', 'v3', 'v8', 'v15']] +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', +} class MSR_Normalization(MetaLayer): @@ -36,10 +57,14 @@ def __init__(self, mode="ALL", **kwargs): raise ValueError(f"Mode {mode} not accepted for sum rules") self.indices = [] + self.divisor_indices = [] if self._msr_enabled: - self.indices += MSR_INDICES + self.indices += [IDX[c] for c in MSR_COMPONENTS] + self.divisor_indices += [IDX[MSR_DENOMINATORS[c]] for c in MSR_COMPONENTS] if self._vsr_enabled: - self.indices += VSR_INDICES + self.divisor_indices += [IDX[VSR_DENOMINATORS[c]] for c in VSR_COMPONENTS] + self.indices += [IDX[c] for c in VSR_COMPONENTS] + self.vsr_factors = [VSR_CONSTANTS[c] for c in VSR_COMPONENTS] # Need this extra dimension for the scatter_to_one operation self.indices = [[i] for i in self.indices] @@ -70,15 +95,13 @@ def call(self, pdf_integrated): norm_constants = [] if self._msr_enabled: - n_ag = [(1.0 - y[IDX['sigma']] - y[IDX['photon']]) / y[IDX['g']]] - norm_constants += n_ag + norm_constants += [(1.0 - y[IDX['sigma']] - y[IDX['photon']])] if self._vsr_enabled: - n_av = [3.0 / y[IDX['v']]] * 3 - n_av3 = [1.0 / y[IDX['v3']]] - n_av8 = [3.0 / y[IDX['v8']]] - n_av15 = [3.0 / y[IDX['v15']]] - norm_constants += n_av + n_av3 + n_av8 + n_av15 + norm_constants += self.vsr_factors + + divisors = op.gather(y, self.divisor_indices, axis=0) + norm_constants = norm_constants / divisors # Fill in the rest of the flavours with 1 norm_constants = op.scatter_to_one( From eb73f5a3d2a38472c865cc03da0d1f0b53729655 Mon Sep 17 00:00:00 2001 From: Aron Date: Tue, 8 Aug 2023 15:21:20 +0200 Subject: [PATCH 17/21] Undo joining of photon integral --- n3fit/src/n3fit/layers/msr_normalization.py | 7 +++++-- n3fit/src/n3fit/msr.py | 11 +++-------- n3fit/src/n3fit/tests/test_msr.py | 16 +++++++++++----- 3 files changed, 19 insertions(+), 15 deletions(-) diff --git a/n3fit/src/n3fit/layers/msr_normalization.py b/n3fit/src/n3fit/layers/msr_normalization.py index ab2d108511..5e1b2a52cf 100644 --- a/n3fit/src/n3fit/layers/msr_normalization.py +++ b/n3fit/src/n3fit/layers/msr_normalization.py @@ -70,7 +70,7 @@ def __init__(self, mode="ALL", **kwargs): super().__init__(**kwargs) - def call(self, pdf_integrated): + def call(self, pdf_integrated, photon_integral): """ Computes the normalization factors for the PDFs: A_g = (1-sigma-photon)/g @@ -85,6 +85,8 @@ def call(self, pdf_integrated): ---------- pdf_integrated: (Tensor(1, 14)) the integrated PDF + photon_integral: (Tensor(1, 1)) + the integrated photon PDF Returns ------- @@ -92,10 +94,11 @@ def call(self, pdf_integrated): The normalization factors per flavour. """ y = pdf_integrated[0] # get rid of the batch dimension + photon_integral = photon_integral[0] # get rid of the batch dimension norm_constants = [] if self._msr_enabled: - norm_constants += [(1.0 - y[IDX['sigma']] - y[IDX['photon']])] + norm_constants += [(1.0 - y[IDX['sigma']] - photon_integral[0])] if self._vsr_enabled: norm_constants += self.vsr_factors diff --git a/n3fit/src/n3fit/msr.py b/n3fit/src/n3fit/msr.py index 474ec7854a..fda3048815 100644 --- a/n3fit/src/n3fit/msr.py +++ b/n3fit/src/n3fit/msr.py @@ -88,15 +88,10 @@ def generate_msr_model_and_grid( # 5. THe input for the photon integral, will be set to 0 if no photons photon_integral = Input(shape=(1,), batch_size=1, name='photon_integral') - # 5b. Insert the photon integral as the first component of the pdf integrals - pdf_integrated = Lambda( - lambda pdf_photon: op.concatenate([pdf_photon[1], pdf_photon[0][:, 1:]], axis=1), - name="join_photon", - )([pdf_integrated, photon_integral]) - # 6. Compute the normalization factor - # For now set the photon component to None - normalization_factor = MSR_Normalization(mode, name="msr_weights")(pdf_integrated) + normalization_factor = MSR_Normalization(mode, name="msr_weights")( + pdf_integrated, photon_integral + ) # 7. Apply the normalization factor to the pdf pdf_normalized = Lambda(lambda pdf_norm: pdf_norm[0] * pdf_norm[1], name="pdf_normalized")( diff --git a/n3fit/src/n3fit/tests/test_msr.py b/n3fit/src/n3fit/tests/test_msr.py index 25547ec22f..9c72f860b5 100644 --- a/n3fit/src/n3fit/tests/test_msr.py +++ b/n3fit/src/n3fit/tests/test_msr.py @@ -8,13 +8,9 @@ def apply_layer_to_fixed_input(layer): np.random.seed(422) pdf_integrated = op.numpy_to_tensor(np.random.normal(size=(1, 14))) - # done this way to stay consistent with the original code photon_integral = op.numpy_to_tensor([np.random.normal(size=(1,))]) - pdf_integrated = op.concatenate( - [op.numpy_to_tensor(photon_integral), pdf_integrated[:, 1:]], axis=1 - ) - return layer(pdf_integrated) + return layer(pdf_integrated, photon_integral) def test_all(): @@ -87,3 +83,13 @@ def test_vsr(): ] ) np.testing.assert_allclose(output, known_output, rtol=1e-5) + + +def main(): + test_all() + test_msr() + test_vsr() + + +if __name__ == '__main__': + main() From 96b798bb82c56b24b89b6f2f4cf1e2d576ebe39f Mon Sep 17 00:00:00 2001 From: Aron Date: Wed, 9 Aug 2023 10:53:53 +0200 Subject: [PATCH 18/21] turn lists into tensors --- n3fit/src/n3fit/layers/msr_normalization.py | 9 ++++++--- n3fit/src/n3fit/tests/test_msr.py | 10 ---------- 2 files changed, 6 insertions(+), 13 deletions(-) diff --git a/n3fit/src/n3fit/layers/msr_normalization.py b/n3fit/src/n3fit/layers/msr_normalization.py index 5e1b2a52cf..c185fb9910 100644 --- a/n3fit/src/n3fit/layers/msr_normalization.py +++ b/n3fit/src/n3fit/layers/msr_normalization.py @@ -64,7 +64,7 @@ def __init__(self, mode="ALL", **kwargs): if self._vsr_enabled: self.divisor_indices += [IDX[VSR_DENOMINATORS[c]] for c in VSR_COMPONENTS] self.indices += [IDX[c] for c in VSR_COMPONENTS] - self.vsr_factors = [VSR_CONSTANTS[c] for c in VSR_COMPONENTS] + self.vsr_factors = op.numpy_to_tensor([VSR_CONSTANTS[c] for c in VSR_COMPONENTS]) # Need this extra dimension for the scatter_to_one operation self.indices = [[i] for i in self.indices] @@ -98,10 +98,13 @@ def call(self, pdf_integrated, photon_integral): norm_constants = [] if self._msr_enabled: - norm_constants += [(1.0 - y[IDX['sigma']] - photon_integral[0])] + norm_constants += [ + op.batchit(1.0 - y[IDX['sigma']] - photon_integral[0], batch_dimension=0) + ] if self._vsr_enabled: - norm_constants += self.vsr_factors + norm_constants += [self.vsr_factors] + norm_constants = op.concatenate(norm_constants, axis=0) divisors = op.gather(y, self.divisor_indices, axis=0) norm_constants = norm_constants / divisors diff --git a/n3fit/src/n3fit/tests/test_msr.py b/n3fit/src/n3fit/tests/test_msr.py index 9c72f860b5..668ba15be7 100644 --- a/n3fit/src/n3fit/tests/test_msr.py +++ b/n3fit/src/n3fit/tests/test_msr.py @@ -83,13 +83,3 @@ def test_vsr(): ] ) np.testing.assert_allclose(output, known_output, rtol=1e-5) - - -def main(): - test_all() - test_msr() - test_vsr() - - -if __name__ == '__main__': - main() From b52a6f4fcf642ae42f62c5a49e5c026c657db955 Mon Sep 17 00:00:00 2001 From: Aron Jansen Date: Thu, 10 Aug 2023 16:30:25 +0200 Subject: [PATCH 19/21] Add module docstring Co-authored-by: Juan M. Cruz-Martinez --- n3fit/src/n3fit/layers/msr_normalization.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/n3fit/src/n3fit/layers/msr_normalization.py b/n3fit/src/n3fit/layers/msr_normalization.py index c185fb9910..86e4b47863 100644 --- a/n3fit/src/n3fit/layers/msr_normalization.py +++ b/n3fit/src/n3fit/layers/msr_normalization.py @@ -1,3 +1,11 @@ +""" + Definition of the imposition of the Momentum Sum Rule and Valence Sum Rules to in the PDF fit. + + In the module level constants ``{MSR/VSR}_COMPONENTS`` the flavours affected by the MSR and VSR are defined. + For the Valence Sum Rule instead `VSR_DENOMINATOR` defines the integral of which flavour are used + to compute the normalization. Note that for a Nf=4 fit `v35=v24=v`. + If the number of flavours were to be changed in the future, this would need to be updated accordingly. +""" from n3fit.backends import MetaLayer from n3fit.backends import operations as op From 09a99b8b841a0a9cf1df311e7efe239a2a7e599c Mon Sep 17 00:00:00 2001 From: Aron Date: Fri, 11 Aug 2023 09:48:56 +0200 Subject: [PATCH 20/21] Avoid duplicate names --- n3fit/src/n3fit/layers/msr_normalization.py | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/n3fit/src/n3fit/layers/msr_normalization.py b/n3fit/src/n3fit/layers/msr_normalization.py index 86e4b47863..4ad312ae85 100644 --- a/n3fit/src/n3fit/layers/msr_normalization.py +++ b/n3fit/src/n3fit/layers/msr_normalization.py @@ -64,17 +64,17 @@ def __init__(self, mode="ALL", **kwargs): else: raise ValueError(f"Mode {mode} not accepted for sum rules") - self.indices = [] + indices = [] self.divisor_indices = [] if self._msr_enabled: - self.indices += [IDX[c] for c in MSR_COMPONENTS] + indices += [IDX[c] for c in MSR_COMPONENTS] self.divisor_indices += [IDX[MSR_DENOMINATORS[c]] for c in MSR_COMPONENTS] if self._vsr_enabled: self.divisor_indices += [IDX[VSR_DENOMINATORS[c]] for c in VSR_COMPONENTS] - self.indices += [IDX[c] for c in VSR_COMPONENTS] + indices += [IDX[c] for c in VSR_COMPONENTS] self.vsr_factors = op.numpy_to_tensor([VSR_CONSTANTS[c] for c in VSR_COMPONENTS]) # Need this extra dimension for the scatter_to_one operation - self.indices = [[i] for i in self.indices] + self.indices = [[i] for i in indices] super().__init__(**kwargs) @@ -103,23 +103,21 @@ def call(self, pdf_integrated, photon_integral): """ y = pdf_integrated[0] # get rid of the batch dimension photon_integral = photon_integral[0] # get rid of the batch dimension - norm_constants = [] + numerators = [] if self._msr_enabled: - norm_constants += [ + numerators += [ op.batchit(1.0 - y[IDX['sigma']] - photon_integral[0], batch_dimension=0) ] - if self._vsr_enabled: - norm_constants += [self.vsr_factors] - norm_constants = op.concatenate(norm_constants, axis=0) + numerators += [self.vsr_factors] + numerators = op.concatenate(numerators, axis=0) divisors = op.gather(y, self.divisor_indices, axis=0) - norm_constants = norm_constants / divisors # Fill in the rest of the flavours with 1 norm_constants = op.scatter_to_one( - norm_constants, indices=self.indices, output_shape=y.shape + numerators / divisors, indices=self.indices, output_shape=y.shape ) return norm_constants From a85d31492cebcc84a74c563cd8f9f99080557a51 Mon Sep 17 00:00:00 2001 From: Aron Date: Fri, 11 Aug 2023 09:50:41 +0200 Subject: [PATCH 21/21] formatting --- n3fit/src/n3fit/layers/msr_normalization.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/n3fit/src/n3fit/layers/msr_normalization.py b/n3fit/src/n3fit/layers/msr_normalization.py index 4ad312ae85..f543676e6c 100644 --- a/n3fit/src/n3fit/layers/msr_normalization.py +++ b/n3fit/src/n3fit/layers/msr_normalization.py @@ -1,6 +1,6 @@ """ Definition of the imposition of the Momentum Sum Rule and Valence Sum Rules to in the PDF fit. - + In the module level constants ``{MSR/VSR}_COMPONENTS`` the flavours affected by the MSR and VSR are defined. For the Valence Sum Rule instead `VSR_DENOMINATOR` defines the integral of which flavour are used to compute the normalization. Note that for a Nf=4 fit `v35=v24=v`. @@ -21,9 +21,7 @@ 'v35': 8, } MSR_COMPONENTS = ['g'] -MSR_DENOMINATORS = { - 'g': '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']