src/utils/neighbors.py

import torch
import src
from src.dependencies.FRNN import frnn
from torch_scatter import scatter
from torch_geometric.utils import coalesce
from src.utils.scatter import scatter_nearest_neighbor


__all__ = [
    'knn_1', 'knn_1_graph', 'knn_2', 'inliers_split', 'outliers_split',
    'inliers_outliers_splits', 'cluster_radius_nn_graph']


def knn_1(
        xyz,
        k,
        r_max=1,
        batch=None,
        oversample=False,
        self_is_neighbor=False,
        verbose=False):
    """Search k-NN for a 3D point cloud xyz. This search differs
    from `knn_2` in that it operates on a single cloud input (search and
    query are the same) and it allows oversampling the neighbors when
    less than `k` neighbors are found within `r_max`. Optionally,
    passing `batch` will ensure the neighbor search does not mix up
    batch items.
    """
    assert isinstance(xyz, torch.Tensor)
    assert k >= 1
    assert xyz.dim() == 2
    assert batch is None or batch.shape[0] == xyz.shape[0]

    # To take the batch into account, we add an offset to the Z
    # coordinates. The offset is designed so that any points from two
    # batch different batch items are separated by at least `r_max + 1`
    batch_offset = 0
    if batch is not None:
        z_offset = xyz[:, 2].max() - xyz[:, 2].min() + r_max + 1
        batch_offset = torch.zeros_like(xyz)
        batch_offset[:, 2] = batch * z_offset

    # Data initialization
    device = xyz.device
    xyz_query = (xyz + batch_offset).view(1, -1, 3)
    xyz_search = (xyz + batch_offset).view(1, -1, 3)
    if not xyz.is_cuda:
        xyz_query = xyz_query.cuda()
        xyz_search = xyz_search.cuda()

    # KNN on GPU. Actual neighbor search now
    k_search = k if self_is_neighbor else k + 1
    distances, neighbors, _, _ = frnn.frnn_grid_points(
        xyz_query, xyz_search, K=k_search, r=r_max)

    # Remove each point from its own neighborhood
    neighbors = neighbors[0] if self_is_neighbor else neighbors[0][:, 1:]
    distances = distances[0] if self_is_neighbor else distances[0][:, 1:]

    # Oversample the neighborhoods where less than k points were found
    if oversample:
        neighbors, distances = oversample_partial_neighborhoods(
            neighbors, distances, k)

    # Restore the neighbors and distances to the input device
    if neighbors.device != device:
        neighbors = neighbors.to(device)
        distances = distances.to(device)

    if not verbose and not src.is_debug_enabled():
        return neighbors, distances

    # Warn the user of partial and empty neighborhoods
    num_nodes = neighbors.shape[0]
    n_missing = (neighbors < 0).sum(dim=1)
    n_partial = (n_missing > 0).sum()
    n_empty = (n_missing == k).sum()
    if n_partial == 0:
        return neighbors, distances

    print(
        f"\nWarning: {n_partial}/{num_nodes} points have partial "
        f"neighborhoods and {n_empty}/{num_nodes} have empty "
        f"neighborhoods (missing neighbors are indicated by -1 indices).")

    return neighbors, distances


def knn_1_graph(
        xyz,
        k,
        r_max=1,
        batch=None,
        oversample=False,
        self_is_neighbor=False,
        verbose=False,
        trim=True):
    """Search k-NN for a 3D point cloud xyz and convert the output into
    torch_geometric's `edge_index`, `edge_attr` format. This search
    differs from `knn_2` in that it operates on a single cloud input
    (search and query are the same) and it allows oversampling the
    neighbors when less than `k` neighbors are found within `r_max`.
    Optionally, passing `batch` will ensure the neighbor search does not
    mix up batch items.

    Importantly, the output graph will be coalesced: duplicate edges
    will be removed. Besides, if `trim=True`, the graph will be further
    reduced using `to_trimmed()` (see function documentation for more
    information).
    """
    # Nearest neighbor search
    neighbors, distances = knn_1(
        xyz,
        k,
        r_max=r_max,
        batch=batch,
        oversample=oversample,
        self_is_neighbor=self_is_neighbor,
        verbose=verbose)

    # Build the corresponding graph
    num_points = xyz.shape[0]
    source = torch.arange(num_points, device=xyz.device).repeat_interleave(k)
    target = neighbors.flatten()
    edge_index = torch.vstack((source, target))
    distances = distances.flatten()

    # Trim edges where points are missing (i.e. -1 neighbor indices)
    missing_point_edge = edge_index[1] == -1
    edge_index = edge_index[:, ~missing_point_edge]
    distances = distances[~missing_point_edge]

    # Trim the graph. This is required before computing the actual
    # nearest points between all cluster pairs. Since this operation is
    # so costly, we first built on a coarse neighborhood edge_index to
    # alleviate compute and memory cost
    if trim:
        from src.utils import to_trimmed
        edge_index, distances = to_trimmed(
            edge_index, edge_attr=distances, reduce='min')
    # Coalesce edges to remove duplicates
    else:
        edge_index, distances = coalesce(
            edge_index, edge_attr=distances, reduce='min')

    return edge_index, distances


def knn_2(
        x_search,
        x_query,
        k,
        r_max=1,
        batch_search=None,
        batch_query=None):
    """Search k-NN of x_query inside x_search, within radius `r_max`.
    Optionally, passing `batch_search` and `batch_query` will ensure the
    neighbor search does not mix up batch items.
    """
    assert isinstance(x_search, torch.Tensor)
    assert isinstance(x_query, torch.Tensor)
    assert k >= 1
    assert x_search.dim() == 2
    assert x_query.dim() == 2
    assert x_query.shape[1] == x_search.shape[1]
    assert bool(batch_search) == bool(batch_query)
    assert batch_search is None or batch_search.shape[0] == x_search.shape[0]
    assert batch_query is None or batch_query.shape[0] == x_query.shape[0]

    k = torch.tensor([k])
    r_max = torch.tensor([r_max])

    # To take the batch into account, we add an offset to the Z
    # coordinates. The offset is designed so that any points from two
    # batch different batch items are separated by at least `r_max + 1`
    batch_search_offset = 0
    batch_query_offset = 0
    if batch_search is not None:
        hi = max(x_search[:, 2].max(), x_query[:, 2].max())
        lo = min(x_search[:, 2].min(), x_query[:, 2].min())
        z_offset = hi - lo + r_max + 1
        batch_search_offset = torch.zeros_like(x_search)
        batch_search_offset[:, 2] = batch_search * z_offset
        batch_query_offset = torch.zeros_like(x_query)
        batch_query_offset[:, 2] = batch_query * z_offset

    # Data initialization
    device = x_search.device
    xyz_query = (x_query + batch_query_offset).view(1, -1, 3).cuda()
    xyz_search = (x_search + batch_search_offset).view(1, -1, 3).cuda()

    # KNN on GPU. Actual neighbor search now
    distances, neighbors, _, _ = frnn.frnn_grid_points(
        xyz_query, xyz_search, K=k, r=r_max)

    # Remove each point from its own neighborhood
    neighbors = neighbors[0].to(device)
    distances = distances[0].to(device)
    if k == 1:
        neighbors = neighbors[:, 0]
        distances = distances[:, 0]

    return neighbors, distances


def inliers_split(
        xyz_query, xyz_search, k_min, r_max=1, recursive=False, q_in_s=False):
    """Optionally recursive inlier search. The `xyz_query` and
    `xyz_search`. Search for points with less than `k_min` neighbors
    within a radius of `r_max`.

    Since removing outliers may cause some points to become outliers
    themselves, this problem can be tackled with the `recursive` option.
    Note that this recursive search holds no guarantee of reasonable
    convergence as one could design a point cloud for given `k_min` and
    `r_max` whose points would all recursively end up as outliers.
    """
    return inliers_outliers_splits(
        xyz_query, xyz_search, k_min, r_max=r_max, recursive=recursive,
        q_in_s=q_in_s)[0]


def outliers_split(
        xyz_query, xyz_search, k_min, r_max=1, recursive=False, q_in_s=False):
    """Optionally recursive outlier search. The `xyz_query` and
    `xyz_search`. Search for points with less than `k_min` neighbors
    within a radius of `r_max`.

    Since removing outliers may cause some points to become outliers
    themselves, this problem can be tackled with the `recursive` option.
    Note that this recursive search holds no guarantee of reasonable
    convergence as one could design a point cloud for given `k_min` and
    `r_max` whose points would all recursively end up as outliers.
    """
    return inliers_outliers_splits(
        xyz_query, xyz_search, k_min, r_max=r_max, recursive=recursive,
        q_in_s=q_in_s)[1]


def inliers_outliers_splits(
        xyz_query, xyz_search, k_min, r_max=1, recursive=False, q_in_s=False):
    """Optionally recursive outlier search. The `xyz_query` and
    `xyz_search`. Search for points with less than `k_min` neighbors
    within a radius of `r_max`.

    Since removing outliers may cause some points to become outliers
    themselves, this problem can be tackled with the `recursive` option.
    Note that this recursive search holds no guarantee of reasonable
    convergence as one could design a point cloud for given `k_min` and
    `r_max` whose points would all recursively end up as outliers.
    """
    # Data initialization
    device = xyz_query.device
    xyz_query = xyz_query.view(1, -1, 3).cuda()
    xyz_search = xyz_search.view(1, -1, 3).cuda()

    # KNN on GPU. Actual neighbor search now
    neighbors = frnn.frnn_grid_points(
        xyz_query, xyz_search, K=k_min + q_in_s, r=r_max)[1]

    # If the Query points are included in the Search points, remove each
    # point from its own neighborhood
    if q_in_s:
        neighbors = neighbors[0][:, 1:]

    # Get the number of found neighbors for each point. Indeed,
    # depending on the cloud properties and the chosen K and radius,
    # some points may receive "-1" neighbors
    n_found_nn = (neighbors != -1).sum(dim=1)

    # Identify points which have less than k_min neighbor. Those are
    # treated as outliers
    mask_outliers = n_found_nn < k_min
    idx_outliers = torch.where(mask_outliers)[0]
    idx_inliers = torch.where(~mask_outliers)[0]

    # Exit here if not recursively searching for outliers
    if not recursive:
        return idx_outliers.to(device), idx_inliers.to(device)

    # Identify the points affected by the removal of the outliers. Those
    # inliers are potential outliers
    idx_potential = torch.where(
        torch.isin(neighbors[idx_inliers], idx_outliers).any(dim=1))[0]

    # Exit here if there are no potential new outliers among the inliers
    if idx_potential.shape[0] == 0:
        return idx_outliers.to(device), idx_inliers.to(device)

    # Recursively search actual outliers among the potential
    xyz_query_sub = xyz_query[0, idx_inliers[idx_potential]]
    xyz_search_sub = xyz_search[0, idx_inliers]
    idx_outliers_sub, idx_inliers_sub = inliers_outliers_splits(
        xyz_query_sub, xyz_search_sub, k_min, r_max=r_max, recursive=True,
        q_in_s=True)

    # Update the outliers mask
    mask_outliers[idx_inliers[idx_potential][idx_outliers_sub]] = True
    idx_outliers = torch.where(mask_outliers)[0]
    idx_inliers = torch.where(~mask_outliers)[0]

    return idx_outliers.to(device), idx_inliers.to(device)


def oversample_partial_neighborhoods(neighbors, distances, k):
    """Oversample partial neighborhoods with less than k points. Missing
    neighbors are indicated by the "-1" index.

    Remarks
      - Neighbors and distances are assumed to be sorted in order of
      increasing distance
      - All neighbors are assumed to have at least one valid neighbor.
      See `search_outliers` to remove points with not enough neighbors
    """
    # Initialization
    assert neighbors.dim() == distances.dim() == 2
    device = neighbors.device

    # Get the number of found neighbors for each point. Indeed,
    # depending on the cloud properties and the chosen K and radius,
    # some points may receive `-1` neighbors
    n_found_nn = (neighbors != -1).sum(dim=1)

    # Identify points which have more than k_min and less than k
    # neighbors within R. For those, we oversample the neighbors to
    # reach k
    idx_partial = torch.where(n_found_nn < k)[0]
    neighbors_partial = neighbors[idx_partial]
    distances_partial = distances[idx_partial]

    # Since the neighbors are sorted by increasing distance, the missing
    # neighbors will always be the last ones. This helps finding their
    # number and position, for oversampling.

    # *******************************************************************
    # The above statement is actually INCORRECT because the outlier
    # removal may produce "-1" neighbors at unexpected positions. So
    # either we manage to treat this in a clean vectorized way, or we
    # fall back to the 2-searches solution...
    # Honestly, this feels like it is getting out of hand, let's keep
    # things simple, since we are not going to save so much computation
    # time with KNN wrt the partition.
    # *******************************************************************

    # For each missing neighbor, compute the size of the discrete set to
    # oversample from.
    n_valid = n_found_nn[idx_partial].repeat_interleave(
        k - n_found_nn[idx_partial])

    # Compute the oversampling row indices.
    idx_x_sampling = torch.arange(
        neighbors_partial.shape[0], device=device).repeat_interleave(
        k - n_found_nn[idx_partial])

    # Compute the oversampling column indices. The 0.9999 factor is a
    # security to handle the case where torch.rand is to close to 1.0,
    # which would yield incorrect sampling coordinates that would in
    # result in sampling '-1' indices (i.e. all we try to avoid here)
    idx_y_sampling = (n_valid * torch.rand(
        n_valid.shape[0], device=device) * 0.9999).floor().long()

    # Apply the oversampling
    idx_missing = torch.where(neighbors_partial == -1)
    neighbors_partial[idx_missing] = neighbors_partial[
        idx_x_sampling, idx_y_sampling]
    distances_partial[idx_missing] = distances_partial[
        idx_x_sampling, idx_y_sampling]

    # Restore the oversampled neighborhoods with the rest
    neighbors[idx_partial] = neighbors_partial
    distances[idx_partial] = distances_partial

    return neighbors, distances


def cluster_radius_nn_graph(
        x_points,
        idx,
        k_max=100,
        gap=0,
        batch=None,
        trim=True,
        cycles=3,
        chunk_size=100000):
    """Compute the radius neighbors of clusters. Two clusters are
    considered neighbors if 2 of their points are distant of `gap` of
    less.

    The underlying strategy searches the cluster centroids within a
    certain radius, based each cluster's radius and the chosen `gap`.
    This approach is a proxy to avoid the actual computation of all
    pointwise distances.

    :param x_points:
    :param idx:
    :param k_max:
    :param gap:
    :param batch:
        Passing `batch` will ensure the neighbor search does
        not mix up batch items. This batch tensor is a tensor of size
        `num_clusters=idx.max() + 1` indicating which batch item each
        cluster belongs to
    :param trim bool
        If True, the output `edge_index` will be trimmed using
        `to_trimmed`, to save compute and memory
    :param cycles int
        Number of iterations. Starting from a point X in set A, one
        cycle accounts for searching the nearest neighbor, in A, of the
        nearest neighbor of X in set B
    :param chunk_size: int, float
        Allows mitigating memory use when computing the neighbors. If
        `chunk_size > 1`, `edge_index` will be processed into chunks of
        `chunk_size`. If `0 < chunk_size < 1`, then `edge_index` will be
        divided into parts of `edge_index.shape[1] * chunk_size` or less
    :return:
    """
    assert batch is None or batch.shape[0] == idx.max() + 1

    device = x_points.device

    # Roughly estimate the diameter and center of each segment. Note we
    # do not use the centroid (center of mass) but rather the center of
    # the bounding box
    bbox_low = scatter(x_points, idx, dim=0, reduce='min')
    bbox_high = scatter(x_points, idx, dim=0, reduce='max')
    diam = (bbox_high - bbox_low).max(dim=1).values
    center = (bbox_high + bbox_low) / 2

    # Conservative heuristic for the global search radius: we search the
    # segments whose centroids are separated by the largest segment
    # diameter plus the input gap. This approximates the true operation
    # we would like to perform (but which is too costly): searching, for
    # each segment, the segments with at least one point within gap.
    # Obviously, the r_search may produce more neighbors than needed and
    # some subsequent pruning will be needed
    r_search = float(diam.max() + gap)
    neighbors, distances = knn_1(center, k_max, r_max=r_search, batch=batch)

    # Build the corresponding edge_index
    num_clusters = idx.max() + 1
    source = torch.arange(num_clusters, device=device).repeat_interleave(k_max)
    target = neighbors.flatten()
    edge_index = torch.vstack((source, target))
    distances = distances.flatten()

    # Trim edges based on the actual segment radii and not the
    # overly-conservative maximum radius used for the search. For this
    # step, we use a gap of `sqrt(3) * gap` to account for some edge
    # case where two 3D boxes touch each other by their corners. This
    # avoids removing neighbors too aggressively before the next step
    # TODO: for some reason, this trimming does not remove as many
    #  neighbors as I'd thought, so I need to call
    #  `scatter_nearest_neighbor` afterward. There must be something
    #  wrong in the `r_max_edge` or the `distances` here. Solving this
    #  is no priority but could maybe avoid the call to
    #  `scatter_nearest_neighbor`
    r_segment = diam / 2
    r_max_edge = r_segment[edge_index].sum(dim=0) + 1.732 * gap
    in_gap_range = distances <= r_max_edge
    edge_index = edge_index[:, in_gap_range]
    distances = distances[in_gap_range]

    # Trim edges where points are missing (i.e. -1 neighbor indices)
    missing_point_edge = edge_index[1] == -1
    edge_index = edge_index[:, ~missing_point_edge]
    distances = distances[~missing_point_edge]

    # Trim the graph. This is required before computing the actual
    # nearest points between all cluster pairs. Since this operation is
    # so costly, we first built on a coarse neighborhood edge_index to
    # alleviate compute and memory cost
    if trim:
        from src.utils import to_trimmed
        edge_index, distances = to_trimmed(
            edge_index, edge_attr=distances, reduce='min')
    # Coalesce edges to remove duplicates
    else:
        edge_index, distances = coalesce(
            edge_index, edge_attr=distances, reduce='min')

    # For each cluster pair in edge_index, compute (approximately) the
    # two closest points (coined "anchors" here). The heuristic used
    # here to find those points runs in O(E) with E the number of
    # edges, which is O(N) with N the number of points. This is a
    # workaround for the actual anchor points search, which is O(N²)
    # TODO: scatter_nearest_neighbor is the bottleneck of cluster_nn_radius(),
    #  we could accelerate things by randomly sampling in the clusters
    anchors = scatter_nearest_neighbor(
        x_points, idx, edge_index, cycles=cycles, chunk_size=chunk_size)[1]
    d_nn = (x_points[anchors[0]] - x_points[anchors[1]]).norm(dim=1)

    # Trim edges wrt the anchor points distance
    in_gap_range = d_nn <= gap
    edge_index = edge_index[:, in_gap_range]
    distances = d_nn[in_gap_range]

    return edge_index, distances