Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor rotations #1780

Merged
merged 4 commits into from
Aug 8, 2023
Merged

Refactor rotations #1780

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
4e0e226
Write rotations specifying rotation axis
APJansen Jul 18, 2023
becbf9d
Add regression test for rotation
APJansen Jul 18, 2023
f830a74
Rename moveaxis to swapaxes
APJansen Jul 27, 2023
9c475e1
Merge branch 'master' into refactor_rotations
APJansen Jul 27, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions n3fit/src/n3fit/backends/keras_backend/operations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -355,6 +355,22 @@ def op_subtract(inputs, **kwargs):
return keras_subtract(inputs, **kwargs)


def swapaxes(tensor, source, destination):
"""
Moves the axis of the tensor from source to destination, as in numpy.swapaxes.
see full `docs <https://numpy.org/doc/stable/reference/generated/numpy.swapaxes.html>`_
"""
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)

APJansen marked this conversation as resolved.
Show resolved Hide resolved

@tf.function
def backend_function(fun_name, *args, **kwargs):
"""
Expand Down
52 changes: 28 additions & 24 deletions n3fit/src/n3fit/layers/rotations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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.swapaxes(rotated, -1, self.rotation_axis)


class FlavourToEvolution(Rotation):
Expand All @@ -53,7 +54,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):
Expand All @@ -64,29 +65,32 @@ 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()
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"""
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

Expand Down
51 changes: 51 additions & 0 deletions n3fit/src/n3fit/tests/test_rotations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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)