1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
|
import random
from typing import Tuple
from torch_geometric.data import Data
from torch_geometric.data.datapipes import functional_transform
from torch_geometric.transforms import BaseTransform
@functional_transform('random_scale')
class RandomScale(BaseTransform):
r"""Scales node positions by a randomly sampled factor :math:`s` within a
given interval, *e.g.*, resulting in the transformation matrix
(functional name: :obj:`random_scale`).
.. math::
\begin{bmatrix}
s & 0 & 0 \\
0 & s & 0 \\
0 & 0 & s \\
\end{bmatrix}
for three-dimensional positions.
Args:
scales (tuple): scaling factor interval, e.g. :obj:`(a, b)`, then scale
is randomly sampled from the range
:math:`a \leq \mathrm{scale} \leq b`.
"""
def __init__(self, scales: Tuple[float, float]) -> None:
assert isinstance(scales, (tuple, list)) and len(scales) == 2
self.scales = scales
def forward(self, data: Data) -> Data:
assert data.pos is not None
scale = random.uniform(*self.scales)
data.pos = data.pos * scale
return data
def __repr__(self) -> str:
return f'{self.__class__.__name__}({self.scales})'
|
