def test_radius(dtype, device): x = tensor([ [-1, -1], [-1, +1], [+1, +1], [+1, -1], [-1, -1], [-1, +1], [+1, +1], [+1, -1], ], dtype, device) y = tensor([ [0, 0], [0, 1], ], dtype, device) batch_x = tensor([0, 0, 0, 0, 1, 1, 1, 1], torch.long, device) batch_y = tensor([0, 1], torch.long, device) edge_index = radius(x, y, 2, max_num_neighbors=4) assert to_set(edge_index) == set([(0, 0), (0, 1), (0, 2), (0, 3), (1, 1), (1, 2), (1, 5), (1, 6)]) edge_index = radius(x, y, 2, batch_x, batch_y, max_num_neighbors=4) assert to_set(edge_index) == set([(0, 0), (0, 1), (0, 2), (0, 3), (1, 5), (1, 6)])
def test(): from torch_cluster import radius from e3nn.math import soft_one_hot_linspace conv = Convolution( irreps_node_input='0e + 1e', irreps_node_output='0e + 1e', irreps_node_attr_input='2x0e', irreps_node_attr_output='3x0e', irreps_edge_attr='0e + 1e', num_edge_scalar_attr=4, radial_layers=1, radial_neurons=50, num_neighbors=3.0, ) pos_in = torch.randn(5, 3) pos_out = torch.randn(2, 3) node_input = torch.randn(5, 4) node_attr_input = torch.randn(5, 2) node_attr_output = torch.randn(2, 3) edge_src, edge_dst = radius(pos_out, pos_in, r=2.0) edge_vec = pos_in[edge_src] - pos_out[edge_dst] edge_attr = o3.spherical_harmonics([0, 1], edge_vec, True) edge_scalar_attr = soft_one_hot_linspace(x=edge_vec.norm(dim=1), start=0.0, end=2.0, number=4, basis='smooth_finite', cutoff=True) conv(node_input, node_attr_input, node_attr_output, edge_src, edge_dst, edge_attr, edge_scalar_attr)
def random_crop_3D(P, F, factor): npoints = P.shape[0] n_points_after_crop = np.round(npoints * factor).astype(np.int) points_max = (P.max(axis=0) * 1000).astype(np.int) points_min = (P.min(axis=0) * 1000).astype(np.int) centroid = np.asarray([ np.random.randint(low=points_min[0], high=points_max[0], dtype=int), np.random.randint(low=points_min[1], high=points_max[1], dtype=int), np.random.randint(low=points_min[2], high=points_max[2], dtype=int) ]) centroid = centroid.astype(np.float32) / 1000 rad = 0.1 inc = 0.2 npoints_inside_sphere = 0 x = torch.from_numpy(P) y = torch.from_numpy(centroid).unsqueeze(0) while npoints_inside_sphere < n_points_after_crop: _, crop = torch_cluster.radius(x, y, rad, max_num_neighbors=n_points_after_crop) npoints_inside_sphere = len(crop) rad = np.round(rad + inc, 1) return P[crop.numpy()], F[crop.numpy()]
def _variance_estimator_sparse(r, pos, f, batch_idx): with torch.no_grad(): assign_index = radius(pos, pos, r, batch_x=batch_idx, batch_y=batch_idx) y_idx, x_idx = assign_index # diff = pos[x_idx] - pos[y_idx] # squared_distance = (diff * diff).sum(dim=-1, keepdim=True) # weights = 1.0 / torch.clamp(squared_distance, min=1e-16) grad_f = (f[x_idx] - f[y_idx])**2 y = scatter_add(grad_f, y_idx, dim=0, dim_size=pos.size(0)) return y
def radius(x, y, r, batch_x=None, batch_y=None, max_num_neighbors=32, num_workers=1): r"""Finds for each element in :obj:`y` all points in :obj:`x` within distance :obj:`r`. Args: x (Tensor): Node feature matrix :math:`\mathbf{X} \in \mathbb{R}^{N \times F}`. y (Tensor): Node feature matrix :math:`\mathbf{Y} \in \mathbb{R}^{M \times F}`. r (float): The radius. batch_x (LongTensor, optional): Batch vector :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each node to a specific example. (default: :obj:`None`) batch_y (LongTensor, optional): Batch vector :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^M`, which assigns each node to a specific example. (default: :obj:`None`) max_num_neighbors (int, optional): The maximum number of neighbors to return for each element in :obj:`y`. (default: :obj:`32`) num_workers (int): Number of workers to use for computation. Has no effect in case :obj:`batch_x` or :obj:`batch_y` is not :obj:`None`, or the input lies on the GPU. (default: :obj:`1`) :rtype: :class:`LongTensor` .. code-block:: python import torch from torch_geometric.nn import radius x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]]) batch_x = torch.tensor([0, 0, 0, 0]) y = torch.Tensor([[-1, 0], [1, 0]]) batch_y = torch.tensor([0, 0]) assign_index = radius(x, y, 1.5, batch_x, batch_y) """ if torch_cluster is None: raise ImportError('`radius` requires `torch-cluster`.') return torch_cluster.radius(x, y, r, batch_x, batch_y, max_num_neighbors, num_workers)
def test_radius(dtype, device): x = tensor([ [-1, -1], [-1, +1], [+1, +1], [+1, -1], [-1, -1], [-1, +1], [+1, +1], [+1, -1], ], dtype, device) y = tensor([ [0, 0], [0, 1], ], dtype, device) batch_x = tensor([0, 0, 0, 0, 1, 1, 1, 1], torch.long, device) batch_y = tensor([0, 1], torch.long, device) out = radius(x, y, 2, batch_x, batch_y, max_num_neighbors=4) assert coalesce(out).tolist() == [[0, 0, 0, 0, 1, 1], [0, 1, 2, 3, 5, 6]]