From d5e14f90e42ba860d749cb7bc9f5d8e8f90ca0c3 Mon Sep 17 00:00:00 2001 From: Aron Date: Sat, 2 Dec 2023 14:18:21 +0100 Subject: [PATCH] Move replica axis to front everywhere --- .../n3fit/backends/keras_backend/MetaModel.py | 10 +- .../backends/keras_backend/operations.py | 14 +-- .../src/n3fit/hyper_optimization/penalties.py | 13 +-- n3fit/src/n3fit/layers/DIS.py | 58 +++++------ n3fit/src/n3fit/layers/DY.py | 14 +-- n3fit/src/n3fit/layers/msr_normalization.py | 44 +++------ n3fit/src/n3fit/layers/rotations.py | 13 +-- n3fit/src/n3fit/layers/x_operations.py | 4 +- n3fit/src/n3fit/model_gen.py | 14 +-- n3fit/src/n3fit/model_trainer.py | 2 +- n3fit/src/n3fit/msr.py | 8 +- n3fit/src/n3fit/stopping.py | 2 +- n3fit/src/n3fit/tests/test_layers.py | 37 +++---- n3fit/src/n3fit/tests/test_msr.py | 98 ++++++++----------- n3fit/src/n3fit/tests/test_multireplica.py | 2 +- n3fit/src/n3fit/tests/test_rotations.py | 4 +- n3fit/src/n3fit/tests/test_xops.py | 4 +- validphys2/src/validphys/photon/compute.py | 4 +- .../validphys/tests/photon/test_compute.py | 16 +-- 19 files changed, 160 insertions(+), 201 deletions(-) diff --git a/n3fit/src/n3fit/backends/keras_backend/MetaModel.py b/n3fit/src/n3fit/backends/keras_backend/MetaModel.py index 9ed23b0823..dfefd38401 100644 --- a/n3fit/src/n3fit/backends/keras_backend/MetaModel.py +++ b/n3fit/src/n3fit/backends/keras_backend/MetaModel.py @@ -396,8 +396,7 @@ def split_replicas(self): if self.single_replica_generator is None: raise ValueError("Trying to generate single replica models with no generator set.") replicas = [] - num_replicas = self.output.shape[-1] - for i_replica in range(num_replicas): + for i_replica in range(self.num_replicas): replica = self.single_replica_generator() replica.set_replica_weights(self.get_replica_weights(i_replica)) @@ -410,14 +409,17 @@ def split_replicas(self): return replicas + @property + def num_replicas(self): + return self.output.shape[1] + def load_identical_replicas(self, model_file): """ From a single replica model, load the same weights into all replicas. """ weights = self._format_weights_from_file(model_file) - num_replicas = self.output.shape[-1] - for i_replica in range(num_replicas): + for i_replica in range(self.num_replicas): self.set_replica_weights(weights, i_replica) def _format_weights_from_file(self, model_file): diff --git a/n3fit/src/n3fit/backends/keras_backend/operations.py b/n3fit/src/n3fit/backends/keras_backend/operations.py index 4ec1ecea4f..f965152f42 100644 --- a/n3fit/src/n3fit/backends/keras_backend/operations.py +++ b/n3fit/src/n3fit/backends/keras_backend/operations.py @@ -261,7 +261,7 @@ def pdf_masked_convolution(raw_pdf, basis_mask): Parameters ---------- pdf: tf.tensor - rank 4 (batchsize, xgrid, flavours, replicas) + rank 4 (batchsize, replicas, xgrid, flavours) basis_mask: tf.tensor rank 2 tensor (flavours, flavours) mask to apply to the pdf convolution @@ -269,18 +269,18 @@ def pdf_masked_convolution(raw_pdf, basis_mask): Return ------ pdf_x_pdf: tf.tensor - rank3 (len(mask_true), xgrid, xgrid, replicas) + rank3 (replicas, len(mask_true), xgrid, xgrid) """ - if raw_pdf.shape[-1] == 1: # only one replica! - pdf = tf.squeeze(raw_pdf, axis=(0, -1)) + if raw_pdf.shape[1] == 1: # only one replica! + pdf = tf.squeeze(raw_pdf, axis=(0, 1)) luminosity = tensor_product(pdf, pdf, axes=0) lumi_tmp = K.permute_dimensions(luminosity, (3, 1, 2, 0)) - pdf_x_pdf = batchit(boolean_mask(lumi_tmp, basis_mask), -1) + pdf_x_pdf = batchit(boolean_mask(lumi_tmp, basis_mask), 0) else: pdf = tf.squeeze(raw_pdf, axis=0) # remove the batchsize - luminosity = tf.einsum('air,bjr->jibar', pdf, pdf) + luminosity = tf.einsum('rai,rbj->rjiba', pdf, pdf) # (xgrid, flavour, xgrid, flavour) - pdf_x_pdf = boolean_mask(luminosity, basis_mask) + pdf_x_pdf = boolean_mask(luminosity, basis_mask, axis=1) return pdf_x_pdf diff --git a/n3fit/src/n3fit/hyper_optimization/penalties.py b/n3fit/src/n3fit/hyper_optimization/penalties.py index 414758ce38..77a156cd3a 100644 --- a/n3fit/src/n3fit/hyper_optimization/penalties.py +++ b/n3fit/src/n3fit/hyper_optimization/penalties.py @@ -53,17 +53,18 @@ def saturation(pdf_model=None, n=100, min_x=1e-6, max_x=1e-4, flavors=None, **_k if flavors is None: flavors = [1, 2] x = np.logspace(np.log10(min_x), np.log10(max_x), n) - x = np.expand_dims(x, axis=[0, -1]) extra_loss = 0.0 - y = pdf_model.predict({"pdf_input": x}) - xpdf = y[0, :, flavors] + x_input = np.expand_dims(x, axis=[0, -1]) + y = pdf_model.predict({"pdf_input": x_input}) + xpdf = y[0, :, :, flavors] # this is now of shape (flavors, replicas, xgrid) - delta_logx = np.diff(np.log10(x), axis=1) - delta_xpdf = np.diff(xpdf, axis=1) + x = np.expand_dims(x, axis=[0, 1]) + delta_logx = np.diff(np.log10(x), axis=2) + delta_xpdf = np.diff(xpdf, axis=2) slope = delta_xpdf / delta_logx - pen = abs(np.mean(slope, axis=1)) + np.std(slope, axis=1) + pen = abs(np.mean(slope, axis=2)) + np.std(slope, axis=2) # sum over flavors # Add a small offset to avoid ZeroDivisionError diff --git a/n3fit/src/n3fit/layers/DIS.py b/n3fit/src/n3fit/layers/DIS.py index 70dec828f6..0aeb2d00e6 100644 --- a/n3fit/src/n3fit/layers/DIS.py +++ b/n3fit/src/n3fit/layers/DIS.py @@ -7,29 +7,31 @@ """ import numpy as np -from .observable import Observable + from n3fit.backends import operations as op +from .observable import Observable + 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, replicas, xgrid, flavours) """ 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,21 +43,21 @@ 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, replicas, xgrid, flavours) - 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: @@ -65,13 +67,13 @@ def call(self, pdf): # Separate the two possible paths this layer can take if self.many_masks: for mask, fktable in zip(self.all_masks, self.fktables): - pdf_masked = op.boolean_mask(pdf, mask, axis=2) - res = op.tensor_product(pdf_masked, fktable, axes=[(1, 2), (2, 1)]) + pdf_masked = op.boolean_mask(pdf, mask, axis=3) + res = op.tensor_product(pdf_masked, fktable, axes=[(2, 3), (2, 1)]) results.append(res) else: - pdf_masked = op.boolean_mask(pdf, self.all_masks[0], axis=2) + pdf_masked = op.boolean_mask(pdf, self.all_masks[0], axis=3) for fktable in self.fktables: - res = op.tensor_product(pdf_masked, fktable, axes=[(1, 2), (2, 1)]) + res = op.tensor_product(pdf_masked, fktable, axes=[(2, 3), (2, 1)]) results.append(res) return self.operation(results) diff --git a/n3fit/src/n3fit/layers/DY.py b/n3fit/src/n3fit/layers/DY.py index dc5fe858da..3b42a394bd 100644 --- a/n3fit/src/n3fit/layers/DY.py +++ b/n3fit/src/n3fit/layers/DY.py @@ -1,7 +1,9 @@ import numpy as np -from .observable import Observable + from n3fit.backends import operations as op +from .observable import Observable + class DY(Observable): """ @@ -30,7 +32,7 @@ def call(self, pdf_raw): Parameters ---------- pdf_in: tensor - rank 4 tensor (batchsize, xgrid, flavours, replicas) + rank 4 tensor (batchsize, replicas, xgrid, flavours) Returns ------- @@ -43,20 +45,20 @@ def call(self, pdf_raw): results = [] if self.many_masks: if self.splitting: - splitted_pdf = op.split(pdf_raw, self.splitting, axis=1) + splitted_pdf = op.split(pdf_raw, self.splitting, axis=2) for mask, pdf, fk in zip(self.all_masks, splitted_pdf, self.fktables): pdf_x_pdf = op.pdf_masked_convolution(pdf, mask) - res = op.tensor_product(fk, pdf_x_pdf, axes=3) + res = op.tensor_product(fk, pdf_x_pdf, axes=[(1, 2, 3), (1, 2, 3)]) results.append(res) else: for mask, fk in zip(self.all_masks, self.fktables): pdf_x_pdf = op.pdf_masked_convolution(pdf_raw, mask) - res = op.tensor_product(fk, pdf_x_pdf, axes=3) + res = op.tensor_product(fk, pdf_x_pdf, axes=[(1, 2, 3), (1, 2, 3)]) results.append(res) else: pdf_x_pdf = op.pdf_masked_convolution(pdf_raw, self.all_masks[0]) for fk in self.fktables: - res = op.tensor_product(fk, pdf_x_pdf, axes=3) + res = op.tensor_product(fk, pdf_x_pdf, axes=[(1, 2, 3), (1, 2, 3)]) results.append(res) # the masked convolution removes the batch dimension diff --git a/n3fit/src/n3fit/layers/msr_normalization.py b/n3fit/src/n3fit/layers/msr_normalization.py index 0711eee86d..0755cad0b4 100644 --- a/n3fit/src/n3fit/layers/msr_normalization.py +++ b/n3fit/src/n3fit/layers/msr_normalization.py @@ -11,38 +11,14 @@ from n3fit.backends import MetaLayer from n3fit.backends import operations as op -IDX = { - 'photon': 0, - 'sigma': 1, - 'g': 2, - 'v': 3, - 'v3': 4, - 'v8': 5, - 'v15': 6, - 'v24': 7, - 'v35': 8, -} +IDX = {'photon': 0, 'sigma': 1, 'g': 2, 'v': 3, 'v3': 4, 'v8': 5, 'v15': 6, 'v24': 7, 'v35': 8} 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'} class MSR_Normalization(MetaLayer): @@ -93,18 +69,20 @@ def call(self, pdf_integrated, photon_integral): Parameters ---------- - pdf_integrated: (Tensor(1, 14, replicas)) + pdf_integrated: (Tensor(1, replicas, 14)) the integrated PDF - photon_integral: (Tensor(1, 1, replicas)) + photon_integral: (Tensor(1, replicas, 1)) the integrated photon PDF Returns ------- - normalization_factor: Tensor(14, replicas) + normalization_factor: Tensor(replicas, 1, 14) 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 + # get rid of batch dimension and put replicas last + reshape = lambda x: op.transpose(x[0]) + y = reshape(pdf_integrated) + photon_integral = reshape(photon_integral) numerators = [] if self._msr_enabled: @@ -122,4 +100,4 @@ def call(self, pdf_integrated, photon_integral): numerators / divisors, indices=self.indices, output_shape=y.shape ) - return norm_constants + return op.batchit(op.transpose(norm_constants), batch_dimension=1) diff --git a/n3fit/src/n3fit/layers/rotations.py b/n3fit/src/n3fit/layers/rotations.py index 4d9e4eaf8f..4ff9272682 100644 --- a/n3fit/src/n3fit/layers/rotations.py +++ b/n3fit/src/n3fit/layers/rotations.py @@ -22,7 +22,7 @@ class Rotation(MetaLayer): rotation_axis of input to be rotated """ - def __init__(self, rotation_matrix, rotation_axis=2, **kwargs): + def __init__(self, rotation_matrix, rotation_axis=3, **kwargs): self.rotation_matrix = op.numpy_to_tensor(rotation_matrix) self.rotation_axis = rotation_axis super().__init__(**kwargs) @@ -47,12 +47,7 @@ class FlavourToEvolution(Rotation): the evolution basis. """ - def __init__( - self, - flav_info, - fitbasis, - **kwargs, - ): + def __init__(self, flav_info, fitbasis, **kwargs): rotation_matrix = pdfbases.fitbasis_to_NN31IC(flav_info, fitbasis) super().__init__(rotation_matrix, **kwargs) @@ -111,7 +106,7 @@ def __init__(self, photons, **kwargs): super().__init__(**kwargs) def register_photon(self, xgrid): - """Compute the photon array of shape (1, xgrid, 1, replicas) and set the layer to be rebuilt""" + """Compute the photon array of shape (1, replicas, xgrid, 1) and set the layer to be rebuilt""" if self._photons_generator: self._pdf_ph = self._photons_generator(xgrid) self.built = False @@ -119,7 +114,7 @@ def register_photon(self, xgrid): def call(self, pdfs): if self._pdf_ph is None: return pdfs - return op.concatenate([self._pdf_ph, pdfs[:, :, 1:]], axis=2) + return op.concatenate([self._pdf_ph, pdfs[:, :, :, 1:]], axis=3) class ObsRotation(MetaLayer): diff --git a/n3fit/src/n3fit/layers/x_operations.py b/n3fit/src/n3fit/layers/x_operations.py index 1d22b78cca..a55b56d2ba 100644 --- a/n3fit/src/n3fit/layers/x_operations.py +++ b/n3fit/src/n3fit/layers/x_operations.py @@ -65,11 +65,11 @@ class xIntegrator(MetaLayer): ---------- grid_weights: np.array weights of the grid - x_axis: int (default=1) + x_axis: int (default=2) axis of the input tensor that corresponds to the x-grid """ - def __init__(self, grid_weights, x_axis=1, **kwargs): + def __init__(self, grid_weights, x_axis=2, **kwargs): self.x_axis = x_axis self.grid_weights = op.flatten(op.numpy_to_tensor(grid_weights)) super().__init__(**kwargs) diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py index 87e66753d9..a700a10fa9 100644 --- a/n3fit/src/n3fit/model_gen.py +++ b/n3fit/src/n3fit/model_gen.py @@ -78,7 +78,7 @@ def _generate_experimental_layer(self, pdf): splitting_layer = op.as_layer( op.split, op_args=[self.dataset_xsizes], - op_kwargs={"axis": 1}, + op_kwargs={"axis": 2}, name=f"{self.name}_split", ) sp_pdf = splitting_layer(pdf) @@ -578,7 +578,7 @@ def pdfNN_layer_generator( Returns ------- pdf_model: n3fit.backends.MetaModel - a model f(x) = y where x is a tensor (1, xgrid, 1) and y a tensor (1, xgrid, out, num_replicas) + a model f(x) = y where x is a tensor (1, xgrid, 1) and y a tensor (1, replicas, xgrid, out) """ # Parse the input configuration if seed is None: @@ -700,13 +700,13 @@ def pdfNN_layer_generator( # Apply NN layers for all replicas to a given input grid def neural_network_replicas(x, postfix=""): - NNs_x = Lambda(lambda nns: op.stack(nns, axis=-1), name=f"NNs{postfix}")( + NNs_x = Lambda(lambda nns: op.stack(nns, axis=1), name=f"NNs{postfix}")( [nn(x) for nn in nn_replicas] ) if subtract_one: x_eq_1_processed = process_input(layer_x_eq_1) - NNs_x_1 = Lambda(lambda nns: op.stack(nns, axis=-1), name=f"NNs{postfix}_x_1")( + NNs_x_1 = Lambda(lambda nns: op.stack(nns, axis=1), name=f"NNs{postfix}_x_1")( [nn(x_eq_1_processed) for nn in nn_replicas] ) NNs_x = subtract_one_layer([NNs_x, NNs_x_1]) @@ -715,7 +715,7 @@ def neural_network_replicas(x, postfix=""): # Apply preprocessing factors for all replicas to a given input grid def preprocessing_replicas(x, postfix=""): - return Lambda(lambda pfs: op.stack(pfs, axis=-1), name=f"prefactors{postfix}")( + return Lambda(lambda pfs: op.stack(pfs, axis=1), name=f"prefactors{postfix}")( [pf(x) for pf in preprocessing_factor_replicas] ) @@ -752,7 +752,7 @@ def compute_unnormalized_pdf(x, postfix=""): # add batch and flavor dimensions photon_integrals = op.batchit(op.batchit(photons.integral)) else: - photon_integrals = np.zeros((1, 1, num_replicas)) + photon_integrals = np.zeros((1, num_replicas, 1)) PDFs_normalized = sumrule_layer( { @@ -773,7 +773,7 @@ def compute_unnormalized_pdf(x, postfix=""): if replica_axis: pdf_model = MetaModel(model_input, PDFs, name=f"PDFs", scaler=scaler) else: - pdf_model = MetaModel(model_input, PDFs[..., 0], name=f"PDFs", scaler=scaler) + pdf_model = MetaModel(model_input, PDFs[:, 0], name=f"PDFs", scaler=scaler) return pdf_model diff --git a/n3fit/src/n3fit/model_trainer.py b/n3fit/src/n3fit/model_trainer.py index 956cd3a0ea..81af8337ab 100644 --- a/n3fit/src/n3fit/model_trainer.py +++ b/n3fit/src/n3fit/model_trainer.py @@ -356,7 +356,7 @@ def _xgrid_generation(self): # The PDF model will be called with a concatenation of all inputs # now the output needs to be splitted so that each experiment takes its corresponding input sp_ar = [[i.shape[1] for i in inputs_unique]] - sp_kw = {"axis": 1} + sp_kw = {"axis": 2} sp_layer = op.as_layer(op.split, op_args=sp_ar, op_kwargs=sp_kw, name="pdf_split") return InputInfo(input_layer, sp_layer, inputs_idx) diff --git a/n3fit/src/n3fit/msr.py b/n3fit/src/n3fit/msr.py index 2c3de11ca6..e781a672d2 100644 --- a/n3fit/src/n3fit/msr.py +++ b/n3fit/src/n3fit/msr.py @@ -53,9 +53,9 @@ def generate_msr_model_and_grid( - input: the input layer of the integration grid """ # 0. Prepare input layers to MSR model - pdf_x = Input(shape=(None, output_dim, replicas), batch_size=1, name="pdf_x") + pdf_x = Input(shape=(replicas, None, output_dim), batch_size=1, name="pdf_x") pdf_xgrid_integration = Input( - shape=(nx, output_dim, replicas), batch_size=1, name="pdf_xgrid_integration" + shape=(replicas, nx, output_dim), batch_size=1, name="pdf_xgrid_integration" ) # 1. Generate the grid and weights that will be used to integrate @@ -81,14 +81,14 @@ def generate_msr_model_and_grid( # 3. Prepare the pdf for integration by dividing by x pdf_integrand = Lambda( - lambda x_pdf: op.batchit(x_pdf[0], batch_dimension=-1) * x_pdf[1], name="pdf_integrand" + lambda x_pdf: op.batchit(x_pdf[0], batch_dimension=1) * x_pdf[1], name="pdf_integrand" )([x_divided, pdf_xgrid_integration]) # 4. Integrate the pdf 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=(1, replicas), batch_size=1, name='photon_integral') + photon_integral = Input(shape=(replicas, 1), batch_size=1, name='photon_integral') # 6. Compute the normalization factor normalization_factor = MSR_Normalization(mode, replicas, name="msr_weights")( diff --git a/n3fit/src/n3fit/stopping.py b/n3fit/src/n3fit/stopping.py index 5488078a29..8884e7f01b 100644 --- a/n3fit/src/n3fit/stopping.py +++ b/n3fit/src/n3fit/stopping.py @@ -349,7 +349,7 @@ def __init__( self._positivity = Positivity(threshold_positivity, pos_sets) # Initialize internal variables for the stopping - self._n_replicas = pdf_model.output_shape[-1] + self._n_replicas = pdf_model.num_replicas self._threshold_chi2 = threshold_chi2 self._stopping_degrees = np.zeros(self._n_replicas, dtype=int) self._counts = np.zeros(self._n_replicas, dtype=int) diff --git a/n3fit/src/n3fit/tests/test_layers.py b/n3fit/src/n3fit/tests/test_layers.py index b8cbb844a0..1cae6c6c6a 100644 --- a/n3fit/src/n3fit/tests/test_layers.py +++ b/n3fit/src/n3fit/tests/test_layers.py @@ -3,11 +3,12 @@ This module checks that the layers do what they would do with numpy """ import dataclasses + import numpy as np -from validphys.pdfbases import fitbasis_to_NN31IC + from n3fit.backends import operations as op import n3fit.layers as layers - +from validphys.pdfbases import fitbasis_to_NN31IC FLAVS = 3 XSIZE = 4 @@ -145,7 +146,7 @@ def test_DIS(): fks = [i.fktable for i in fktables] obs_layer = layers.DIS(fktables, fks, ope, nfl=FLAVS) pdf = np.random.rand(XSIZE, FLAVS) - kp = op.numpy_to_tensor(np.expand_dims(pdf, 0)) + kp = op.numpy_to_tensor([[pdf]]) # add batch and replica dimension # generate the n3fit results result_tensor = obs_layer(kp) result = op.evaluate(result_tensor) @@ -169,8 +170,7 @@ def test_DY(): fks = [i.fktable for i in fktables] obs_layer = layers.DY(fktables, fks, ope, nfl=FLAVS) pdf = np.random.rand(XSIZE, FLAVS) - # Add batch dimension (0) and replica dimension (-1) - kp = op.numpy_to_tensor(np.expand_dims(pdf, [0, -1])) + kp = op.numpy_to_tensor([[pdf]]) # add batch and replica dimension # generate the n3fit results result_tensor = obs_layer(kp) result = op.evaluate(result_tensor) @@ -201,15 +201,15 @@ def test_rotation_flavour(): {"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) + pdf = np.ones(8) # Vector in the flavour basis v_i + pdf = np.expand_dims(pdf, axis=[0, 1, 2]) # Add batch, replica, x dimensions mat = fitbasis_to_NN31IC(flav_info, "FLAVOUR") # 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 + res_np = np.tensordot(pdf, mat, (3, 0)) # Vector in the evolution basis u_j=R_ij*vi # Apply the rotation through the rotation layer - x = op.numpy_to_tensor(x) + pdf = op.numpy_to_tensor(pdf) rotmat = layers.FlavourToEvolution(flav_info, "FLAVOUR") - res_layer = rotmat(x) + res_layer = rotmat(pdf) assert np.alltrue(res_np == res_layer) @@ -226,15 +226,15 @@ def test_rotation_evol(): {"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) + pdf = np.ones(8) # Vector in the flavour basis v_i + pdf = np.expand_dims(pdf, axis=[0, 1, 2]) # Add batch, replica, x dimensions mat = fitbasis_to_NN31IC(flav_info, "EVOL") # 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 + res_np = np.tensordot(pdf, mat, (3, 0)) # Vector in the evolution basis u_j=R_ij*vi # Apply the rotation through the rotation layer - x = op.numpy_to_tensor(x) + pdf = op.numpy_to_tensor(pdf) rotmat = layers.FlavourToEvolution(flav_info, "EVOL") - res_layer = rotmat(x) + res_layer = rotmat(pdf) assert np.alltrue(res_np == res_layer) @@ -260,6 +260,7 @@ def test_mask(): ret = masker(fi) np.testing.assert_allclose(ret, masked_fi * rn_val, rtol=1e-5) + def test_addphoton_init(): """Test AddPhoton class.""" addphoton = layers.AddPhoton(photons=None) @@ -268,13 +269,15 @@ def test_addphoton_init(): np.testing.assert_equal(addphoton._photons_generator, 1234) np.testing.assert_equal(addphoton._pdf_ph, None) -class FakePhoton(): + +class FakePhoton: def __call__(self, xgrid): return [np.exp(-xgrid)] + def test_compute_photon(): photon = FakePhoton() addphoton = layers.AddPhoton(photons=photon) - xgrid = np.geomspace(1e-4, 1., 10) + xgrid = np.geomspace(1e-4, 1.0, 10) addphoton.register_photon(xgrid) np.testing.assert_allclose(addphoton._pdf_ph, [np.exp(-xgrid)]) diff --git a/n3fit/src/n3fit/tests/test_msr.py b/n3fit/src/n3fit/tests/test_msr.py index e51076741c..9f76d3f421 100644 --- a/n3fit/src/n3fit/tests/test_msr.py +++ b/n3fit/src/n3fit/tests/test_msr.py @@ -6,9 +6,9 @@ def apply_layer_to_fixed_input(layer): np.random.seed(422) - pdf_integrated = op.numpy_to_tensor(np.random.normal(size=(1, 14,1))) + pdf_integrated = op.numpy_to_tensor(np.random.normal(size=(1, 1, 14))) - photon_integral = op.numpy_to_tensor([np.random.normal(size=(1,1))]) + photon_integral = op.numpy_to_tensor([np.random.normal(size=(1, 1))]) return layer(pdf_integrated, photon_integral) @@ -17,25 +17,26 @@ def test_all(): layer = MSR_Normalization(mode='ALL') output = apply_layer_to_fixed_input(layer) known_output = op.numpy_to_tensor( - np.expand_dims( + [ [ - 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, - ], - -1, - ) + [ + 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) @@ -44,25 +45,7 @@ def test_msr(): layer = MSR_Normalization(mode='MSR') output = apply_layer_to_fixed_input(layer) known_output = op.numpy_to_tensor( - np.expand_dims( - [ - 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, - ], - -1, - ) + [[[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) @@ -71,24 +54,25 @@ def test_vsr(): layer = MSR_Normalization(mode='VSR') output = apply_layer_to_fixed_input(layer) known_output = op.numpy_to_tensor( - np.expand_dims( + [ [ - 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, - ], - -1, - ) + [ + 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) diff --git a/n3fit/src/n3fit/tests/test_multireplica.py b/n3fit/src/n3fit/tests/test_multireplica.py index d187560a4e..c3191bef6d 100644 --- a/n3fit/src/n3fit/tests/test_multireplica.py +++ b/n3fit/src/n3fit/tests/test_multireplica.py @@ -6,7 +6,7 @@ def test_replica_split(): """Check that multi replica pdf and concatenated single output pdfs agree""" num_replicas = 3 - replica_axis = -1 + replica_axis = 1 fake_fl = [ {"fl": i, "largex": [0, 1], "smallx": [1, 2]} for i in ["u", "ubar", "d", "dbar", "c", "g", "s", "sbar"] diff --git a/n3fit/src/n3fit/tests/test_rotations.py b/n3fit/src/n3fit/tests/test_rotations.py index a639fde249..d10d9153b9 100644 --- a/n3fit/src/n3fit/tests/test_rotations.py +++ b/n3fit/src/n3fit/tests/test_rotations.py @@ -8,8 +8,10 @@ def test_fk(): rotation = FkRotation() gridpoints = 2 np.random.seed(0) - pdf = op.numpy_to_tensor(np.random.rand(1, gridpoints, 9)) + pdf = op.numpy_to_tensor(np.random.rand(1, 1, gridpoints, 9)) pdf_rotated = rotation(pdf) + # extract single replica + pdf_rotated = pdf_rotated[:, 0] pdf_rotated_known = op.numpy_to_tensor( [ [ diff --git a/n3fit/src/n3fit/tests/test_xops.py b/n3fit/src/n3fit/tests/test_xops.py index db030ae350..957f573f06 100644 --- a/n3fit/src/n3fit/tests/test_xops.py +++ b/n3fit/src/n3fit/tests/test_xops.py @@ -38,10 +38,10 @@ def test_xdivide_indices(): def test_xintegrator(): np.random.seed(42) weights = np.random.rand(5, 1) - pdf = op.numpy_to_tensor(np.random.rand(1, 5, 8)) + pdf = op.numpy_to_tensor(np.random.rand(1, 1, 5, 8)) xint = xIntegrator(weights) xint_out = xint(pdf) xint_out_reference = np.array( - [[0.405455, 0.878931, 0.937715, 0.906214, 1.984154, 1.147975, 1.642387, 1.549858]] + [[[0.405455, 0.878931, 0.937715, 0.906214, 1.984154, 1.147975, 1.642387, 1.549858]]] ) np.testing.assert_allclose(xint_out.numpy(), xint_out_reference, rtol=1e-05) diff --git a/validphys2/src/validphys/photon/compute.py b/validphys2/src/validphys/photon/compute.py index 79a90b8c77..56872540be 100644 --- a/validphys2/src/validphys/photon/compute.py +++ b/validphys2/src/validphys/photon/compute.py @@ -187,14 +187,14 @@ def __call__(self, xgrid): Returns ------- photon values : nd.array - array of photon values with shape (1,xgrid,1,replicas) + array of photon values with shape (1, replicas, xgrid, 1) """ return np.stack( [ self.interpolator[id](xgrid[0, :, 0])[np.newaxis, :, np.newaxis] for id in range(len(self.replicas)) ], - axis=-1, + axis=1, ) @property diff --git a/validphys2/src/validphys/tests/photon/test_compute.py b/validphys2/src/validphys/tests/photon/test_compute.py index 0a204f00fd..0d4c7ea57d 100644 --- a/validphys2/src/validphys/tests/photon/test_compute.py +++ b/validphys2/src/validphys/tests/photon/test_compute.py @@ -203,18 +203,8 @@ def test_betas(): """test betas for different nf""" test_theory = API.theoryid(theoryid=THEORY_QED) alpha = Alpha(test_theory.get_description(), 1e8) - vec_beta0 = [ - -0.5305164769729844, - -0.6719875374991137, - -0.7073553026306458, - -0.8488263631567751, - ] - vec_b1 = [ - 0.17507043740108488, - 0.1605510390839295, - 0.1538497783221655, - 0.1458920311675707, - ] + vec_beta0 = [-0.5305164769729844, -0.6719875374991137, -0.7073553026306458, -0.8488263631567751] + vec_b1 = [0.17507043740108488, 0.1605510390839295, 0.1538497783221655, 0.1458920311675707] for nf in range(3, 6 + 1): np.testing.assert_allclose(alpha.betas_qed[nf][0], vec_beta0[nf - 3], rtol=1e-7) np.testing.assert_allclose(alpha.betas_qed[nf][1], vec_b1[nf - 3], rtol=1e-7) @@ -288,6 +278,6 @@ def test_photon(): photon_Q0 = pdfs_final[ph_id] photon_fiatlux = XGRID * photon_Q0 - photon_validphys = photon(XGRID[np.newaxis, :, np.newaxis])[0, :, 0, 0] + photon_validphys = photon(XGRID[np.newaxis, :, np.newaxis])[0, 0, :, 0] np.testing.assert_allclose(photon_fiatlux, photon_validphys, rtol=1e-7)