neighbors

Gaussian kernel density estimation over packed batches.

Functions:

gaussian_kernel_density –

Computes the Gaussian Kernel Density (KDE) for a given tensor.

gaussian_kernel_density ¶

gaussian_kernel_density(
    x: Tensor, batch: Tensor, bandwidth: float
) -> Tensor

Computes the Gaussian Kernel Density (KDE) for a given tensor.

Parameters:

x (Tensor) –

The input tensor of shape \((N, C)\).
batch (Tensor) –

The batch tensor of shape \((N,)\).
bandwidth (float) –

The bandwidth of the Gaussian kernel.

Returns:

Tensor –

The Gaussian Kernel Density of shape \((N,)\).

Example

import torch
from torch_pointcloud.utils.neighbors import gaussian_kernel_density

pos = torch.randn(100, 3)
batch = torch.zeros(100, dtype=torch.long)
bandwidth = 0.1

density = gaussian_kernel_density(pos, batch, bandwidth)
print(density.shape)
# torch.Size([100])