def clusterize(α, x, scale=None, labels=None):
"""
Performs a simple 'voxelgrid' clustering on the input measure,
putting points into cubic bins of size 'scale' = σ_c.
The weights are summed, and the centroid position is that of the bin's center of mass.
Most importantly, the "fine" lists of weights and points are *sorted*
so that clusters are *contiguous in memory*: this allows us to perform
kernel truncation efficiently on the GPU.
If
[α_c, α], [x_c, x], [x_ranges] = clusterize(α, x, σ_c),
then
α_c[k], x_c[k] correspond to
α[x_ranges[k,0]:x_ranges[k,1]], x[x_ranges[k,0]:x_ranges[k,1],:]
"""
if labels is None and scale is None: # No clustering, single-scale Sinkhorn on the way...
return [α], [x], []
else: # As of today, only two-scale Sinkhorn is implemented:
# Compute simple (voxel-like) class labels:
x_lab = grid_cluster(x, scale) if labels is None else labels
# Compute centroids and weights:
ranges_x, x_c, α_c = cluster_ranges_centroids(x, x_lab, weights=α)
# Make clusters contiguous in memory:
(α, x), x_labels = sort_clusters((α, x), x_lab)
return [α_c, α], [x_c, x], [ranges_x]
def _Gauss_block_sparse_pre(self, x: torch.tensor, y: torch.tensor, K_ij: LazyTensor):
'''
Helper function to preprocess data for block-sparse reduction
of the Gaussian kernel
Args:
x[np.array], y[np.array] = arrays giving rise to Gaussian kernel K(x,y)
K_ij[LazyTensor_n] = symbolic representation of K(x,y)
eps[float] = size for square bins
Returns:
K_ij[LazyTensor_n] = symbolic representation of K(x,y) with
set sparse ranges
'''
# labels for low dimensions
if x.shape[1] < 4 or y.shape[1] < 4:
x_labels = grid_cluster(x, self.eps)
y_labels = grid_cluster(y, self.eps)
# range and centroid per class
x_ranges, x_centroids, _ = cluster_ranges_centroids(x, x_labels)
y_ranges, y_centroids, _ = cluster_ranges_centroids(y, y_labels)
else:
# labels for higher dimensions
x_labels, x_centroids = self._KMeans(x)
y_labels, y_centroids = self._KMeans(y)
# compute ranges
x_ranges = cluster_ranges(x_labels)
y_ranges = cluster_ranges(y_labels)
# sort points
x, x_labels = sort_clusters(x, x_labels)
y, y_labels = sort_clusters(y, y_labels)
# Compute a coarse Boolean mask:
D = torch.sum((x_centroids[:, None, :] - y_centroids[None, :, :]) ** 2, 2)
keep = D < (self.mask_radius) ** 2
# mask -> set of integer tensors
ranges_ij = from_matrix(x_ranges, y_ranges, keep)
K_ij.ranges = ranges_ij # block-sparsity pattern
return K_ij
def kneighbors(self, y):
if use_cuda:
torch.cuda.synchronize()
d = ((y.unsqueeze(1) - self.c.unsqueeze(0))**2).sum(-1)
y_labels = torch.argmin(d, dim=1)
y_ranges, _, _ = cluster_ranges_centroids(y, self.cl)
y, y_labels = sort_clusters(y, y_labels)
ranges_ij = from_matrix(self.x_ranges, y_ranges, self.keep)
y_LT = LazyTensor(y.unsqueeze(0))
D_ij = ((y_LT - self.x)**2).sum(-1)
D_ij.ranges = ranges_ij
return D_ij.argKmin(K=self.k, axis=1)
def fit(self, x, use_torch=True, clusters=50, a=5):
cl, c = KMeans(x, clusters)
self.c = c
#update cluster assignment
if use_torch:
d = ((x.unsqueeze(1) - c.unsqueeze(0))**2).sum(-1)
self.cl = torch.argmin(d, dim=1)
else:
self.cl = k_argmin(x, c)
if use_cuda:
torch.cuda.synchronize()
#get KNN graph for the clusters
if use_torch:
self.ncl = k_argmin_torch(c, c, k=a)
else:
c1 = LazyTensor(c.unsqueeze(1))
c2 = LazyTensor(c.unsqueeze(0))
d = ((c1 - c2)**2).sum(-1)
self.ncl = d.argKmin(K=a, dim=1) #get a nearest clusters
#get the ranges and centroids
self.x_ranges, _, _ = cluster_ranges_centroids(x, self.cl)
#
x, x_labels = sort_clusters(x, self.cl) #sort dataset to match ranges
self.x = LazyTensor(x.unsqueeze(1)) #store dataset
r = torch.arange(clusters).repeat(a, 1).T.reshape(-1).long()
self.keep = torch.zeros([clusters, clusters], dtype=torch.bool)
self.keep[r, self.ncl.flatten()] = True
return self
def sorted(self, x, labels=None):
if labels is None:
labels = self.cl
x, _ = sort_clusters(x, labels)
return x
start = time.time()
x_ranges, x_centroids, _ = cluster_ranges_centroids(x, x_labels)
y_ranges, y_centroids, _ = cluster_ranges_centroids(y, y_labels)
if use_cuda:
torch.cuda.synchronize()
end = time.time()
print("Compute ranges+centroids : {:.4f}s".format(end - start))
###############################################
# Finally, we can **sort** our points according to their
# labels, making sure that **all clusters are stored contiguously in memory**:
from pykeops.torch.cluster import sort_clusters
start = time.time()
x, x_labels = sort_clusters(x, x_labels)
y, y_labels = sort_clusters(y, y_labels)
if use_cuda:
torch.cuda.synchronize()
end = time.time()
print("Sort the points : {:.4f}s".format(end - start))
####################################################################
# Cluster-Cluster binary mask
# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
#
# The key idea behind KeOps's block-sparsity mode
# is that as soon as data points are sorted,
# **we can manage the reduction scheme through a small, coarse boolean mask**
# whose values encode whether or not we should perform computations
# at a finer scale.
def kernel_multiscale(α,
x,
β,
y,
blur=.05,
kernel=None,
name=None,
truncate=5,
diameter=None,
cluster_scale=None,
verbose=False,
**kwargs):
if truncate is None or name == "energy":
return kernel_online(α,
x,
β,
y,
blur=blur,
kernel=kernel,
truncate=truncate,
name=name,
**kwargs)
# Renormalize our point cloud so that blur = 1:
kernel, x, y = kernel_preprocess(kernel, name, x, y, blur)
# Don't forget to normalize the diameter too!
if cluster_scale is None:
D = x.shape[-1]
if diameter is None:
diameter = max_diameter(x.view(-1, D), y.view(-1, D))
else:
diameter = diameter / blur
cluster_scale = diameter / (np.sqrt(D) * 2000**(1 / D))
# Put our points in cubic clusters:
cell_diameter = cluster_scale * np.sqrt(x.shape[1])
x_lab = grid_cluster(x, cluster_scale)
y_lab = grid_cluster(y, cluster_scale)
# Compute the ranges and centroids of each cluster:
ranges_x, x_c, α_c = cluster_ranges_centroids(x, x_lab, weights=α)
ranges_y, y_c, β_c = cluster_ranges_centroids(y, y_lab, weights=β)
if verbose:
print("{}x{} clusters, computed at scale = {:2.3f}".format(
len(x_c), len(y_c), cluster_scale))
# Sort the clusters, making them contiguous in memory:
(α, x), x_lab = sort_clusters((α, x), x_lab)
(β, y), y_lab = sort_clusters((β, y), y_lab)
with torch.no_grad(): # Compute our block-sparse reduction ranges:
# Compute pairwise distances between clusters:
C_xx = squared_distances(x_c, x_c)
C_yy = squared_distances(y_c, y_c)
C_xy = squared_distances(x_c, y_c)
# Compute the boolean masks:
keep_xx = (C_xx <= (truncate + cell_diameter)**2)
keep_yy = (C_yy <= (truncate + cell_diameter)**2)
keep_xy = (C_xy <= (truncate + cell_diameter)**2)
# Compute the KeOps reduction ranges:
ranges_xx = from_matrix(ranges_x, ranges_x, keep_xx)
ranges_yy = from_matrix(ranges_y, ranges_y, keep_yy)
ranges_xy = from_matrix(ranges_x, ranges_y, keep_xy)
return kernel_keops(kernel,
α,
x,
β,
y,
ranges_xx=ranges_xx,
ranges_yy=ranges_yy,
ranges_xy=ranges_xy)
X_i[:, :, 0, :] *= gamma
X_i[:, :, -1, :] *= gamma
Y_j[:, :, 0, :] *= gamma
Y_j[:, :, -1, :] *= gamma
###############################################################################
# Optimizing performances
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#
# Contiguous memory accesses are critical for performances on the GPU.
#
from pykeops.torch.cluster import sort_clusters, cluster_ranges
ranges_j = cluster_ranges(labels_j) # Ranges for all clusters
Y_j, labels_j = sort_clusters(
Y_j, labels_j) # Make sure that all clusters are contiguous in memory
C = len(ranges_j) # Number of classes
if C != labels_j.max() + 1:
raise ValueError("???")
for j, (start_j, end_j) in enumerate(ranges_j):
if start_j >= end_j:
raise ValueError(f"The {j}-th cluster of the atlas seems to be empty.")
###############################################################################
# Each fiber is sampled with 20 points in R^3.
# Thus, one tractogram is a matrix of size n x 60 where n is the number of fibers
# The atlas is labelled, wich means that each fiber belong to a cluster.
# This is summarized by the vector labels_j of size n x 1. labels_j[i] is the label of the fiber i.
def kernel_multiscale(α,
x,
β,
y,
blur=0.05,
kernel=None,
name=None,
truncate=5,
diameter=None,
cluster_scale=None,
potentials=False,
verbose=False,
**kwargs):
if truncate is None or name == "energy":
return kernel_online(α.unsqueeze(0),
x.unsqueeze(0),
β.unsqueeze(0),
y.unsqueeze(0),
blur=blur,
kernel=kernel,
truncate=truncate,
name=name,
potentials=potentials,
**kwargs)
# Renormalize our point cloud so that blur = 1:
# Center the point clouds just in case, to prevent numeric overflows:
center = (x.mean(-2, keepdim=True) + y.mean(-2, keepdim=True)) / 2
x, y = x - center, y - center
x_ = x / blur
y_ = y / blur
# Don't forget to normalize the diameter too!
if cluster_scale is None:
D = x.shape[-1]
if diameter is None:
diameter = max_diameter(x_.view(-1, D), y_.view(-1, D))
else:
diameter = diameter / blur
cluster_scale = diameter / (np.sqrt(D) * 2000**(1 / D))
# Put our points in cubic clusters:
cell_diameter = cluster_scale * np.sqrt(x_.shape[-1])
x_lab = grid_cluster(x_, cluster_scale)
y_lab = grid_cluster(y_, cluster_scale)
# Compute the ranges and centroids of each cluster:
ranges_x, x_c, α_c = cluster_ranges_centroids(x_, x_lab, weights=α)
ranges_y, y_c, β_c = cluster_ranges_centroids(y_, y_lab, weights=β)
if verbose:
print("{}x{} clusters, computed at scale = {:2.3f}".format(
len(x_c), len(y_c), cluster_scale))
# Sort the clusters, making them contiguous in memory:
(α, x), x_lab = sort_clusters((α, x), x_lab)
(β, y), y_lab = sort_clusters((β, y), y_lab)
with torch.no_grad(): # Compute our block-sparse reduction ranges:
# Compute pairwise distances between clusters:
C_xx = squared_distances(x_c, x_c)
C_yy = squared_distances(y_c, y_c)
C_xy = squared_distances(x_c, y_c)
# Compute the boolean masks:
keep_xx = C_xx <= (truncate + cell_diameter)**2
keep_yy = C_yy <= (truncate + cell_diameter)**2
keep_xy = C_xy <= (truncate + cell_diameter)**2
# Compute the KeOps reduction ranges:
ranges_xx = from_matrix(ranges_x, ranges_x, keep_xx)
ranges_yy = from_matrix(ranges_y, ranges_y, keep_yy)
ranges_xy = from_matrix(ranges_x, ranges_y, keep_xy)
return kernel_loss(
α,
x,
β,
y,
blur=blur,
kernel=kernel,
name=name,
potentials=potentials,
use_keops=True,
ranges_xx=ranges_xx,
ranges_yy=ranges_yy,
ranges_xy=ranges_xy,
)
