Source code for torch_geometric.utils.geodesic

import torch
import numpy as np

try:
    import gdist
except ImportError:
    gdist = None


[docs]def geodesic_distance(pos, face, src=None, dest=None, norm=True, max_distance=None): r"""Computes (normalized) geodesic distances of a mesh given by :obj:`pos` and :obj:`face`. If :obj:`src` and :obj:`dest` are given, this method only computes the geodesic distances for the respective source and target node-pairs. .. note:: This function requires the :obj:`gdist` package. To install, run :obj:`pip install cython && pip install gdist`. Args: pos (Tensor): The node positions. face (LongTensor): The face indices. src (LongTensor, optional): If given, only compute geodesic distances for the specified source indices. (default: :obj:`None`) dest (LongTensor, optional): If given, only compute geodesic distances for the specified target indices. (default: :obj:`None`) norm (bool, optional): Normalizes geodesic distances by :math:`\sqrt{\textrm{area}(\mathcal{M})}`. (default: :obj:`True`) max_distance (float, optional): If given, only yields results for geodesic distances less than :obj:`max_distance`. This will speed up runtime dramatically. (default: :obj:`None`) :rtype: Tensor """ if gdist is None: raise ImportError('Package `gdist` could not be found.') max_distance = float('inf') if max_distance is None else max_distance if norm: area = (pos[face[1]] - pos[face[0]]).cross(pos[face[2]] - pos[face[0]]) norm = (area.norm(p=2, dim=1) / 2).sum().sqrt().item() else: norm = 1.0 dtype = pos.dtype pos = pos.detach().cpu().to(torch.double).numpy() face = face.detach().t().cpu().to(torch.int).numpy() if src is None and dest is None: out = gdist.local_gdist_matrix(pos, face, max_distance * norm).toarray() / norm return torch.from_numpy(out).to(dtype) if src is None: src = np.arange(pos.shape[0], dtype=np.int32) else: src = src.detach().cpu().to(torch.int).numpy() dest = None if dest is None else dest.detach().cpu().to(torch.int).numpy() outs = [] for i in range(len(src)): s = src[i:i + 1] d = None if dest is None else dest[i:i + 1] out = gdist.compute_gdist(pos, face, s, d, max_distance * norm) / norm out = torch.from_numpy(out).to(dtype) outs.append(out) out = torch.cat(outs, dim=0) if dest is None: out = out.view(-1, pos.shape[0]) return out