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 _kneighbors(self, y): """ Obtain the k nearest neighbors of the query dataset y """ if self.__x is None: raise ValueError( "Input dataset not fitted yet! Call .fit() first!") if self.__device and self.tools.device(y) != self.__device: raise ValueError( "Input dataset and query dataset must be on same device") if len(y.shape) != 2: raise ValueError("Query dataset must be a 2D tensor") if self.__x.shape[-1] != y.shape[-1]: raise ValueError("Query and dataset must have same dimensions") if self.__normalise: y = y / self.tools.repeat(self.tools.norm(y, 2, -1), y.shape[1]).reshape(-1, y.shape[1]) y = self.tools.contiguous(y) y_labels = self.__assign(y) y_ranges, _, _ = cluster_ranges_centroids(y, y_labels) self.__y_ranges = y_ranges y, y_labels = self.__sort_clusters(y, y_labels, store_x=False) x_LT = self.__LazyTensor(self.tools.unsqueeze(self.__x, 0)) y_LT = self.__LazyTensor(self.tools.unsqueeze(y, 1)) D_ij = self.__distance(y_LT, x_LT) ranges_ij = from_matrix(y_ranges, self.__x_ranges, self.__keep) D_ij.ranges = ranges_ij nn = D_ij.argKmin(K=self.__k, axis=1) return self.__unsort(nn)
def kneighbors(self, y): ''' Obtain the k nearest neighbors of the query dataset y ''' if self.__x is None: raise ValueError( 'Input dataset not fitted yet! Call .fit() first!') if type(y) != torch.Tensor: raise ValueError("Query dataset must be a torch tensor") if y.device != self.__device: raise ValueError( 'Input dataset and query dataset must be on same device') if len(y.shape) != 2: raise ValueError('Query dataset must be a 2D tensor') if self.__x.shape[-1] != y.shape[-1]: raise ValueError('Query and dataset must have same dimensions') if use_cuda: torch.cuda.synchronize() y = y.contiguous() y_labels = self.__assign(y) y_ranges, _, _ = cluster_ranges_centroids(y, y_labels) self.__y_ranges = y_ranges y, y_labels = self.__sort_clusters(y, y_labels, store_x=False) x_LT = LazyTensor(self.__x.unsqueeze(0).to(self.__device).contiguous()) y_LT = LazyTensor(y.unsqueeze(1).to(self.__device).contiguous()) D_ij = ((y_LT - x_LT)**2).sum(-1) ranges_ij = from_matrix(y_ranges, self.__x_ranges, self.__keep) D_ij.ranges = ranges_ij nn = D_ij.argKmin(K=self.__k, axis=1) return self.__unsort(nn)
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 _final_brute_force(self, nearest_clusters, query_pts): """ Final brute force search over clusters in cluster method""" if torch.cuda.is_available(): torch.cuda.synchronize() k = self.k x = self.data_orig.to(self.device) x_labels = self.clusters.long() y = query_pts.to(self.device) y_labels = nearest_clusters[:, 0] x = x.contiguous() y = y.contiguous() x_labels = x_labels.to(self.device) y_labels = y_labels.to(self.device) clusters, a = self.graph.shape r = torch.arange(clusters).repeat(a, 1).T.reshape(-1).long() keep = torch.zeros([clusters, clusters], dtype=torch.bool).to(self.device) keep[r, self.graph.flatten()] = True keep += torch.eye(clusters).bool().to(self.device) x_ranges, x_centroids, _ = cluster_ranges_centroids(x, x_labels) y_ranges, y_centroids, _ = cluster_ranges_centroids(y, y_labels) x, x_labels = self.__sort_clusters(x, x_labels, store_x=True) y, y_labels = self.__sort_clusters(y, y_labels, store_x=False) x_LT = LazyTensor(x.unsqueeze(0).to(self.device).contiguous()) y_LT = LazyTensor(y.unsqueeze(1).to(self.device).contiguous()) D_ij = self.distance(y_LT, x_LT) x_ranges = x_ranges.to(self.device) y_ranges = y_ranges.to(self.device) ranges_ij = from_matrix(y_ranges, x_ranges, keep) D_ij.ranges = ranges_ij nn = D_ij.argKmin(K=k, axis=1) return self.__unsort(nn)
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 fit(self, x, clusters=50, a=5, n=15): ''' Fits the main dataset ''' if type(x) != torch.Tensor: raise ValueError('Input must be a torch tensor') if type(clusters) != int: raise ValueError('Clusters must be an integer') if clusters >= len(x): raise ValueError( 'Number of clusters must be less than length of dataset') if type(a) != int: raise ValueError( 'Number of clusters to search over must be an integer') if a > clusters: raise ValueError( 'Number of clusters to search over must be less than total number of clusters' ) if len(x.shape) != 2: raise ValueError('Input must be a 2D array') x = x.contiguous() self.__device = x.device cl, c = self.__KMeans(x, clusters, Niter=n) self.__c = c cl = self.__assign(x) if use_cuda: torch.cuda.synchronize() ncl = self.__k_argmin(c, c, k=a) self.__x_ranges, _, _ = cluster_ranges_centroids(x, cl) x, x_labels = self.__sort_clusters(x, cl, store_x=True) self.__x = x r = torch.arange(clusters).repeat(a, 1).T.reshape(-1).long() self.__keep = torch.zeros([clusters, clusters], dtype=torch.bool).to(self.__device) self.__keep[r, ncl.flatten()] = True return self
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],:] """ perm = None # did we sort the point cloud at some point? Here's the permutation. 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_labels, perm = torch.sort(x_lab.view(-1)) α, x = α[perm], x[perm] # N.B.: the lines above were return to replace a call to # 'sort_clusters' which does not return the permutation, # an information that is needed to de-permute the dual potentials # if they are required by the user. # (α, x), x_labels = sort_clusters( (α,x), x_lab) return [α_c, α], [x_c, x], [ranges_x], perm
x_labels = grid_cluster(x, eps) # class labels y_labels = grid_cluster(y, eps) # class labels if use_cuda: torch.cuda.synchronize() end = time.time() print("Perform clustering : {:.4f}s".format(end - start)) ############################################### # Once (integer) cluster labels have been computed, # we can compute the **centroids** and **memory footprint** of each class: from pykeops.torch.cluster import cluster_ranges_centroids # Compute one range and centroid per class: 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)
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)
a_i.requires_grad = True x_i.requires_grad = True b_j.requires_grad = True # Compute the loss + gradients: Loss_xy = Loss(a_i, x_i, b_j, y_j) [F_i, G_j, dx_i] = grad(Loss_xy, [a_i, b_j, x_i]) # The generalized "Brenier map" is (minus) the gradient of the Sinkhorn loss # with respect to the Wasserstein metric: BrenierMap = -dx_i / (a_i.view(-1, 1) + 1e-7) # Compute the coarse measures for display ---------------------------------- x_lab = grid_cluster(x_i, cluster_scale) _, x_c, a_c = cluster_ranges_centroids(x_i, x_lab, weights=a_i) y_lab = grid_cluster(y_j, cluster_scale) _, y_c, b_c = cluster_ranges_centroids(y_j, y_lab, weights=b_j) # Fancy display: ----------------------------------------------------------- ax = plt.subplot(((Nits - 1) // 3 + 1), 3, i + 1) ax.scatter([10], [10]) # shameless hack to prevent a slight change of axis... display_potential(ax, G_j, "#E2C5C5") display_potential(ax, F_i, "#C8DFF9") if blur > cluster_scale: display_samples(ax, y_j, b_j, [(0.55, 0.55, 0.95, 0.2)]) display_samples(ax, x_i, a_i, [(0.95, 0.55, 0.55, 0.2)], v=BrenierMap)
def _fit( self, x, clusters=50, a=5, Niter=15, device=None, backend=None, approx=False, n=50, ): """ Fits the main dataset """ if type(clusters) != int: raise ValueError("Clusters must be an integer") if clusters >= len(x): raise ValueError( "Number of clusters must be less than length of dataset") if type(a) != int: raise ValueError( "Number of clusters to search over must be an integer") if a > clusters: raise ValueError( "Number of clusters to search over must be less than total number of clusters" ) if len(x.shape) != 2: raise ValueError("Input must be a 2D array") if self.__normalise: x = x / self.tools.repeat(self.tools.norm(x, 2, -1), x.shape[1]).reshape(-1, x.shape[1]) # if we want to use the approximation in Kmeans, and our metric is angular, switch to full angular metric if approx and self.__metric == "angular": self.__update_metric("angular_full") x = self.tools.contiguous(x) self.__device = device self.__backend = backend cl, c = self.tools.kmeans( x, self.__distance, clusters, Niter=Niter, device=self.__device, approx=approx, n=n, ) self.__c = c cl = self.__assign(x) ncl = self.__k_argmin(c, c, k=a) self.__x_ranges, _, _ = cluster_ranges_centroids(x, cl) x, x_labels = self.__sort_clusters(x, cl, store_x=True) self.__x = x r = self.tools.repeat( self.tools.arange(clusters, device=self.__device), a) self.__keep = self.tools.to( self.tools.zeros([clusters, clusters], dtype=bool), self.__device) self.__keep[r, ncl.flatten()] = True return self
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, )