"""Definition of point cloud."""
import gsops.backend as gs
import sklearn.neighbors as neighbors
from geomfum.io import load_pointcloud
from geomfum.metric import HeatDistanceMetric
from geomfum.shape.shape_utils import (
compute_edge_tangent_vectors,
compute_tangent_frames,
)
from ._base import Shape
[docs]
class PointCloud(Shape):
"""Unstructured point cloud with k-NN connectivity and differential operators.
Parameters
----------
vertices : array-like, shape=[n_vertices, 3]
Vertices of the point cloud.
"""
def __init__(self, vertices):
super().__init__(is_mesh=False)
self.vertices = gs.asarray(vertices)
self.n_neighbors = 30
self._knn_graph = None
self._vertex_areas = None
self._vertex_normals = None
self._vertex_tangent_frames = None
self._edges = None
self._edge_tangent_vectors = None
self._dist_matrix = None
self.metric = None
[docs]
@classmethod
def from_file(cls, filename):
"""Load point cloud from file.
Returns
-------
mesh : PointCloud
A point cloud.
"""
vertices = load_pointcloud(filename)
return cls(vertices)
@property
def n_vertices(self):
"""Number of points.
Returns
-------
n_vertices : int
"""
return self.vertices.shape[0]
@property
def knn_graph(self):
"""K-nearest neighbors connectivity graph.
Returns
-------
knn_info : dict
Dictionary containing:
- 'indices': array-like, shape=[n_vertices, k] - neighbor indices for each vertex
- 'distances': array-like, shape=[n_vertices, k] - distances to neighbors
- 'k': int - number of neighbors
- 'nbrs_model': sklearn.neighbors.NearestNeighbors - fitted model for reuse
"""
if self._knn_graph is None:
vertices_np = gs.to_numpy(gs.to_device(self.vertices, "cpu"))
neigs = neighbors.NearestNeighbors(
n_neighbors=self.n_neighbors, algorithm="kd_tree"
).fit(vertices_np)
distances, indices = neigs.kneighbors(vertices_np)
self._knn_graph = {
"indices": indices,
"distances": distances,
"k": self.n_neighbors,
"neighbors_model": neigs,
}
return self._knn_graph
@property
def vertex_normals(self):
"""Normal vectors estimated from local neighborhoods using PCA.
Returns
-------
normals : array-like, shape=[n_vertices, 3]
Normalized per-vertex normals estimated from local neighborhoods using PCA.
"""
if self._vertex_normals is None:
neighbor_indices = gs.array(self.knn_graph["indices"])
all_neighborhoods = self.vertices[neighbor_indices]
centroids = gs.mean(all_neighborhoods, axis=1)
local_neighborhoods = all_neighborhoods - centroids[:, None, :]
cov_matrices = gs.einsum(
"ijk,ijl->ikl", local_neighborhoods, local_neighborhoods
) / (self.n_neighbors - 1)
try:
_, _, v = gs.linalg.svd(cov_matrices)
normals = v[:, :, 2]
except Exception:
normals = gs.zeros_like(self.vertices)
normals[:, 2] = 1.0
# orient normals consistently, if normal is more aligned with inward direction, flip it
neighbor_vectors = all_neighborhoods - self.vertices[:, None, :]
avg_neighbor_direction = gs.mean(neighbor_vectors, axis=1)
dot_products = gs.sum(normals * avg_neighbor_direction, axis=1)
flip_mask = dot_products > 0
normals[flip_mask] *= -1
# Normalize normals
norms = gs.linalg.norm(normals, axis=1, keepdims=True)
normals = normals / (norms + 1e-12)
self._vertex_normals = normals
return self._vertex_normals
@property
def vertex_tangent_frames(self):
"""Local orthonormal coordinate frames at each point.
Returns
-------
tangent_frame : array-like, shape=[n_vertices, 3, 3]
Tangent frame of the mesh, where:
- [n_vertices, 0, :] are the X basis vectors
- [n_vertices, 1, :] are the Y basis vectors
- [n_vertices, 2, :] are the vertex normals
"""
if self._vertex_tangent_frames is None:
self._vertex_tangent_frames = compute_tangent_frames(
self.vertices, self.vertex_normals
)
return self._vertex_tangent_frames
@property
def edges(self):
"""Edge connectivity from k-NN graph."""
if self._edges is None:
neighbor_indices = gs.array(self.knn_graph["indices"])
edge_inds_from = gs.repeat(
gs.arange(self.vertices.shape[0]), self.n_neighbors
)
self._edges = gs.stack((edge_inds_from, neighbor_indices.flatten()))
return self._edges
@property
def edge_tangent_vectors(self):
"""Edge vectors projected onto local tangent planes.
Returns
-------
edge_tangent_vectors : array-like, shape=[n_edges, 2]
Tangent vectors of the edges, projected onto the local tangent plane.
"""
if self._edge_tangent_vectors is None:
edge_tangent_vectors = compute_edge_tangent_vectors(
self.vertices,
self.edges,
self.vertex_tangent_frames,
)
self._edge_tangent_vectors = edge_tangent_vectors
return self._edge_tangent_vectors
@property
def dist_matrix(self):
"""Pairwise distances between all points using the equipped metric.
Returns
-------
_dist_matrix : array-like, shape=[n_vertices, n_vertices]
Metric distance matrix.
"""
if self._dist_matrix is None:
if self.metric is None:
raise ValueError("Metric is not set.")
self._dist_matrix = self.metric.dist_matrix()
return self._dist_matrix
[docs]
def equip_with_metric(self, metric):
"""Equip point cloud with a distance metric.
Parameters
----------
metric : class
A metric class to use for the point cloud.
"""
if metric == HeatDistanceMetric:
self.metric = metric.from_registry(which="pp3d", shape=self)
else:
self.metric = metric(self)
self._dist_matrix = None