Local neighbourhood for sparsity

[alireza78h]

I'm a newbie to KeOps and still trying to get familiar with it, but it seems pretty interesting.

My work is mostly on imaging, and I'm trying to kinda replicate something like a bilateral filter using KeOps.

To put it simply, suppose I have a signal defined on a regular 2D grid, where each sample has a location (i,j). Computing the Gaussian kernel matrix (K_ij) is straightforward, but I want to only keep the entries of (K_ij) that are close to each other in terms of location.

It would be great if I could consider a d×d rectangular neighbourhood around each sample and only keep the values within that neighbourhood (L_infinity norm). For now, though, I think using the L2 distance works fine.

I tried to write a minimal working example here with the help of https://www.kernel-operations.io/keops/python/api/numpy/Cluster.html#pykeops.numpy.cluster.from_matrix

import torch
from pykeops.torch import LazyTensor
from pykeops.torch.cluster import from_matrix
# Suppose I have an NxN 2D regular grid and my D dimensional signal is defined on this grid
N = 64
D = 5

z = torch.randn(N*N, D).cuda()

# Each sample in x has a location (i,j)

i, j = torch.meshgrid(
    torch.arange(N, device=z.device), 
    torch.arange(N, device=z.device),
    indexing='ij'  # Makes i "row-like" and j "column-like"
)

z_i = LazyTensor( z[:,None,:] )  # x_i.shape = (N**2, 1, D)
z_j = LazyTensor( z[None,:,:] )  # y_j.shape = ( 1, N**2,D)

D_ij = ((z_i - z_j)**2).sum(dim=2)  # Symbolic (N**2,N**2,1) matrix of squared distances
K_ij = (- D_ij).exp()               # Symbolic (N**2,N**2,1 Gaussian kernel matrix


# Now I want to sparsify K_ij

dist = (i - j)**2
keep = (dist <= 10)

# I am not sure about this part, am I passing these arguments correct? 

from_matrix(torch.arange(N*N,device=z.device),torch.arange(N*N,device=z.device),keep)

I saw some related posts form @jeanfeydy (thanks for your insight) in issues like #121 and #148 and #24 but couldn't find a solution for that.

One solution is using a mask and multiply it by K_ij, I don't know how to use the mask though as it seems incompatible with int, but I wanted to ask how to change range of K_ij based on the mask to leverage block-sparsity pattern.