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
|
import random
from torch_geometric.data import Data
from torch_geometric.data.datapipes import functional_transform
from torch_geometric.transforms import BaseTransform
@functional_transform('random_flip')
class RandomFlip(BaseTransform):
"""Flips node positions along a given axis randomly with a given
probability (functional name: :obj:`random_flip`).
Args:
axis (int): The axis along the position of nodes being flipped.
p (float, optional): Probability that node positions will be flipped.
(default: :obj:`0.5`)
"""
def __init__(self, axis: int, p: float = 0.5) -> None:
self.axis = axis
self.p = p
def forward(self, data: Data) -> Data:
assert data.pos is not None
if random.random() < self.p:
pos = data.pos.clone()
pos[..., self.axis] = -pos[..., self.axis]
data.pos = pos
return data
def __repr__(self) -> str:
return f'{self.__class__.__name__}(axis={self.axis}, p={self.p})'
|
