From 4e0e22630e79374d48029443fe93fe2d32bbf8f1 Mon Sep 17 00:00:00 2001 From: Aron Date: Tue, 18 Jul 2023 09:50:37 +0200 Subject: [PATCH 1/3] Write rotations specifying rotation axis --- .../n3fit/backends/keras_backend/operations.py | 14 ++++++++++++++ n3fit/src/n3fit/layers/rotations.py | 18 +++++++++--------- 2 files changed, 23 insertions(+), 9 deletions(-) diff --git a/n3fit/src/n3fit/backends/keras_backend/operations.py b/n3fit/src/n3fit/backends/keras_backend/operations.py index efff803b92..b5541ecf9f 100644 --- a/n3fit/src/n3fit/backends/keras_backend/operations.py +++ b/n3fit/src/n3fit/backends/keras_backend/operations.py @@ -354,6 +354,20 @@ def op_subtract(inputs, **kwargs): """ return keras_subtract(inputs, **kwargs) +def moveaxis(tensor, source, destination): + """ + Moves the axis of the tensor from source to destination + """ + indices = list(range(tensor.shape.rank)) + if source < 0: + source += tensor.shape.rank + if destination < 0: + destination += tensor.shape.rank + + indices[source], indices[destination] = indices[destination], indices[source] + + return tf.transpose(tensor, indices) + @tf.function def backend_function(fun_name, *args, **kwargs): diff --git a/n3fit/src/n3fit/layers/rotations.py b/n3fit/src/n3fit/layers/rotations.py index ed58cdbe15..dc61bb2027 100644 --- a/n3fit/src/n3fit/layers/rotations.py +++ b/n3fit/src/n3fit/layers/rotations.py @@ -18,14 +18,13 @@ class Rotation(MetaLayer): ---------- rotation_matrix: np.array rotation matrix - axes: int or list - if given a number, contracts as many indices as given - if given a list (of tuples) contracts indices according to op.tensor_product + rotation_axis: int + rotation_axis of input to be rotated """ - def __init__(self, rotation_matrix, axes=1, **kwargs): + def __init__(self, rotation_matrix, rotation_axis=2, **kwargs): self.rotation_matrix = op.numpy_to_tensor(rotation_matrix) - self.axes = axes + self.rotation_axis = rotation_axis super().__init__(**kwargs) def is_identity(self): @@ -37,7 +36,9 @@ def is_identity(self): return np.allclose(self.rotation_matrix, iden) def call(self, x_raw): - return op.tensor_product(x_raw, self.rotation_matrix, self.axes) + rotated = op.tensor_product(x_raw, self.rotation_matrix, [self.rotation_axis, 0]) + # this puts the rotated axis back in the original place + return op.moveaxis(rotated, -1, self.rotation_axis) class FlavourToEvolution(Rotation): @@ -45,7 +46,6 @@ class FlavourToEvolution(Rotation): Rotates from the flavour basis to the evolution basis. """ - def __init__( self, flav_info, @@ -53,7 +53,7 @@ def __init__( **kwargs, ): rotation_matrix = pdfbases.fitbasis_to_NN31IC(flav_info, fitbasis) - super().__init__(rotation_matrix, axes=1, **kwargs) + super().__init__(rotation_matrix, **kwargs) class FkRotation(Rotation): @@ -67,7 +67,7 @@ class FkRotation(Rotation): def __init__(self, output_dim=14, name="evolution", **kwargs): self.output_dim = output_dim rotation_matrix = self._create_rotation_matrix() - super().__init__(rotation_matrix, axes=1, name=name, **kwargs) + super().__init__(rotation_matrix, name=name, **kwargs) def _create_rotation_matrix(self): """Create the rotation matrix""" From becbf9d0e2d51b0d94bae3366d567281a53b0577 Mon Sep 17 00:00:00 2001 From: Aron Date: Tue, 18 Jul 2023 09:54:02 +0200 Subject: [PATCH 2/3] Add regression test for rotation --- n3fit/src/n3fit/tests/test_rotations.py | 51 +++++++++++++++++++++++++ 1 file changed, 51 insertions(+) create mode 100644 n3fit/src/n3fit/tests/test_rotations.py diff --git a/n3fit/src/n3fit/tests/test_rotations.py b/n3fit/src/n3fit/tests/test_rotations.py new file mode 100644 index 0000000000..a639fde249 --- /dev/null +++ b/n3fit/src/n3fit/tests/test_rotations.py @@ -0,0 +1,51 @@ +import numpy as np + +from n3fit.backends import operations as op +from n3fit.layers import FkRotation, FlavourToEvolution + + +def test_fk(): + rotation = FkRotation() + gridpoints = 2 + np.random.seed(0) + pdf = op.numpy_to_tensor(np.random.rand(1, gridpoints, 9)) + pdf_rotated = rotation(pdf) + pdf_rotated_known = op.numpy_to_tensor( + [ + [ + [ + 0.0, + 0.5488135, + 0.71518934, + 0.60276335, + 0.5448832, + 0.4236548, + 0.96366274, + 0.60276335, + 0.60276335, + 0.6458941, + 0.4375872, + -3.0182784, + 0.5488135, + 0.5488135, + ], + [ + 0.0, + 0.3834415, + 0.79172504, + 0.5288949, + 0.56804454, + 0.92559665, + 0.83261985, + 0.5288949, + 0.5288949, + 0.07103606, + 0.0871293, + 0.30256793, + 0.3834415, + 0.3834415, + ], + ] + ] + ) + np.testing.assert_allclose(pdf_rotated.numpy(), pdf_rotated_known.numpy(), rtol=1e-5) From f830a7446ff4f78addb6ee7c70efd3c92fe88f1e Mon Sep 17 00:00:00 2001 From: Aron Date: Thu, 27 Jul 2023 13:07:35 +0200 Subject: [PATCH 3/3] Rename moveaxis to swapaxes --- .../backends/keras_backend/operations.py | 6 ++- n3fit/src/n3fit/layers/rotations.py | 38 ++++++++++--------- 2 files changed, 25 insertions(+), 19 deletions(-) diff --git a/n3fit/src/n3fit/backends/keras_backend/operations.py b/n3fit/src/n3fit/backends/keras_backend/operations.py index b5541ecf9f..1bd18fa10e 100644 --- a/n3fit/src/n3fit/backends/keras_backend/operations.py +++ b/n3fit/src/n3fit/backends/keras_backend/operations.py @@ -354,9 +354,11 @@ def op_subtract(inputs, **kwargs): """ return keras_subtract(inputs, **kwargs) -def moveaxis(tensor, source, destination): + +def swapaxes(tensor, source, destination): """ - Moves the axis of the tensor from source to destination + Moves the axis of the tensor from source to destination, as in numpy.swapaxes. + see full `docs `_ """ indices = list(range(tensor.shape.rank)) if source < 0: diff --git a/n3fit/src/n3fit/layers/rotations.py b/n3fit/src/n3fit/layers/rotations.py index dc61bb2027..4a970c0472 100644 --- a/n3fit/src/n3fit/layers/rotations.py +++ b/n3fit/src/n3fit/layers/rotations.py @@ -38,7 +38,7 @@ def is_identity(self): def call(self, x_raw): rotated = op.tensor_product(x_raw, self.rotation_matrix, [self.rotation_axis, 0]) # this puts the rotated axis back in the original place - return op.moveaxis(rotated, -1, self.rotation_axis) + return op.swapaxes(rotated, -1, self.rotation_axis) class FlavourToEvolution(Rotation): @@ -46,6 +46,7 @@ class FlavourToEvolution(Rotation): Rotates from the flavour basis to the evolution basis. """ + def __init__( self, flav_info, @@ -64,6 +65,7 @@ class FkRotation(Rotation): The input to this layer is a `pdf_raw` variable which is expected to have a shape (1, None, 9), and it is then rotated to an output (1, None, 14) """ + def __init__(self, output_dim=14, name="evolution", **kwargs): self.output_dim = output_dim rotation_matrix = self._create_rotation_matrix() @@ -71,22 +73,24 @@ def __init__(self, output_dim=14, name="evolution", **kwargs): def _create_rotation_matrix(self): """Create the rotation matrix""" - array = np.array([ - [0, 0, 0, 0, 0, 0, 0, 0, 0], # photon - [1, 0, 0, 0, 0, 0, 0, 0, 0], # sigma - [0, 1, 0, 0, 0, 0, 0, 0, 0], # g - [0, 0, 1, 0, 0, 0, 0, 0, 0], # v - [0, 0, 0, 1, 0, 0, 0, 0, 0], # v3 - [0, 0, 0, 0, 1, 0, 0, 0, 0], # v8 - [0, 0, 0, 0, 0, 0, 0, 0, 1], # v15 - [0, 0, 1, 0, 0, 0, 0, 0, 0], # v24 - [0, 0, 1, 0, 0, 0, 0, 0, 0], # v35 - [0, 0, 0, 0, 0, 1, 0, 0, 0], # t3 - [0, 0, 0, 0, 0, 0, 1, 0, 0], # t8 - [1, 0, 0, 0, 0, 0, 0,-4, 0], # t15 (c-) - [1, 0, 0, 0, 0, 0, 0, 0, 0], # t24 - [1, 0, 0, 0, 0, 0, 0, 0, 0], # t35 - ]) + array = np.array( + [ + [0, 0, 0, 0, 0, 0, 0, 0, 0], # photon + [1, 0, 0, 0, 0, 0, 0, 0, 0], # sigma + [0, 1, 0, 0, 0, 0, 0, 0, 0], # g + [0, 0, 1, 0, 0, 0, 0, 0, 0], # v + [0, 0, 0, 1, 0, 0, 0, 0, 0], # v3 + [0, 0, 0, 0, 1, 0, 0, 0, 0], # v8 + [0, 0, 0, 0, 0, 0, 0, 0, 1], # v15 + [0, 0, 1, 0, 0, 0, 0, 0, 0], # v24 + [0, 0, 1, 0, 0, 0, 0, 0, 0], # v35 + [0, 0, 0, 0, 0, 1, 0, 0, 0], # t3 + [0, 0, 0, 0, 0, 0, 1, 0, 0], # t8 + [1, 0, 0, 0, 0, 0, 0, -4, 0], # t15 (c-) + [1, 0, 0, 0, 0, 0, 0, 0, 0], # t24 + [1, 0, 0, 0, 0, 0, 0, 0, 0], # t35 + ] + ) tensor = op.numpy_to_tensor(array.T) return tensor