How to do Sparse Tensor Indexing with Coordinates Tensor?

[rancheng]

Hi I recently have some tried to do Spatial Propagation on sparse tensor and am wondering if you have implemented the sparse tensor indexing function somewhere so that I can directly use.

Here's a simple example case:

# generate 15 random coordinate entries range from 0 to 3
a_C = torch.empty(15, 2, dtype=torch.long).random_(3)
# generate 15 feature values
a_F = torch.rand(15,1)
# generate sparse tensor
a_S = ME.SparseTensor(a_F, a_C)
# generate random coordinate to do indexing
b_C = torch.empty(2, 2, dtype=torch.long).random_(3)
# I want to select features according to b_C
b_F = a_S[b_C]  # <<<< this is what I want

I want to select the feature according to the coordinates of a given tensor, and those coordinates doesn't have values, it can be set as nan or a customized value.

If you know how to do this, I'm willing to implement the function and start a pull request.

[chrischoy]

Currently, there is no function that exactly does that, but you can make one with get_coords_map in https://github.com/StanfordVL/MinkowskiEngine/blob/master/MinkowskiEngine/MinkowskiCoords.py#L291.

1. create a coords_key of a_S (which is already done when you call ME.SparseTensor)

a_key = a_S.coords_key

2. create a coords_key of b_C by

b_key = a_S.coords_man.create_coords_key(b_C.int())  # it has to be int

3. Find input to output mapping or A to B mapping

a_indices, b_indices = a_S.coords_man.get_coords_map(a_key, b_key)

If there is no match, the lists will be empty. But if you have matches, the order should be

b_F = torch.zeros(len(b_C), a_F.shape[1])  #  placeholder
if len(b_indices) > 0:
  b_F[b_indices, :] = a_F[a_indices, :]

I'll put this into a function in the next cycle.

[rancheng]

Thanks very much for the quick reply! I'll close this issue since it perfectly solved my issue.

[chrischoy]

Hi I just made a function "features_at_coords(query_coords: IntTensor)". Could you see if it works?