Mask indexing when dim(mask)!=dim(source)

1

I am having a system of atoms from which I want to create a graph. Let’s also say for clarity that I am only interested in distances between atoms as feature at the moment.

Moreover, let’s assume we have a function which from input coordinates spits out distances for all atoms to all other atoms within constant radius (in other words it takes only neighboring atoms which are within a sphere of this radius). This is very important because it means that it is very likely that different atoms will have different number of neighbors

Now… that’s when I encountered a problem
So, as I just mentioned I have a function which returns me a 1-D tensor consisting of distances for all atoms (i.e. torch.Size([#DistancesFromAtomsToAllOtherAtoms]), from which I want to build 2-D tensor of the dimensions torch.Size([#NumAtoms, #Neighbors] but the second dimensions “Neighbors” has to be padded.

For example let’s say we are looking at three atoms and the first has 20 neighbors, second 23 and third 21, then 1-D tensor will have the shape of torch.Size([64]), while the 2-D tensor needs to have shape of torch.Size([3,23]) which asks for 69 elements and I get error when I just try to index new tensor with some mask (i.e output[mask] = distances), where mask is boolean tensor with the dimensions torch.Size([3,23])

For this specific example Pytorch error would be:

RuntimeError: The size of tensor a (23) must match the size of tensor b (64) at non-singleton dimension 1

Finally, it’s worth noting that the same process with analogous functions, and with arrays with same dimensions as tensors above, work perfectly in numpy. It just puts zero in resulting tensor whenever the value of mask array is False, and takes value from source array when the value of mask array is True.

Hope someone could suggest me a way to do the same indexing process in Pytorch, while retaining the possibility to do backprop pass on this operation.

Thanks in advance,
Ivan

2

Could you post an example (using random values) for the working numpy code, please?

3

The problem is that it can not be random. It works for a specific case when the number of True values in a mask array is equal to the number of elements in source array. But I will give you an example for my code. Hope it’s not too cumbersome, and the benefit is it’s also physical and easier to understand. I will continue with the example of 3 atoms

Okay four arrays are important here. First one is n_nbh. This array tells me how many neighbors each of the atoms have. It has the dimensions of (number_atoms, max_number_of_neighbors)

array([[16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16],
       [15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15, 15],
       [ 9,  9,  9,  9,  9,  9,  9,  9,  9,  9,  9,  9,  9,  9,  9,  9]])

And here is one example of such array in which first atom has 16, second 15, and third 9 neighbors. Why is it repeating? Well, so that I can make my mask array. However, to create condition I need second array nbh_range

array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15],
       [ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15],
       [ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15]],
      dtype=int32)

This one just gives me range until number of maximum neighbors so that now I can create mask

idx=np.arange(n_atoms) # in our case number of atoms is 3
mask = np.zeros((n_atoms, n_max_nbh), dtype=np.bool_) #n_max_nbh = 16 in our case
mask[idx, :] = nbh_range < n_nbh
mask
array([[ True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True,  True],
       [ True,  True,  True,  True,  True,  True,  True,  True,  True,
         True,  True,  True,  True,  True,  True, False],
       [ True,  True,  True,  True,  True,  True,  True,  True,  True,
        False, False, False, False, False, False, False]])

Furthermore, we have our distances array which is of shape (40,) which comes from 16+15+9 as it signifies distances to all other atoms within a cutoff radius for each atom

dist
array([2.14767316, 2.        , 1.18848643, 2.91247318, 1.        ,
       0.46097722, 2.71661554, 2.66129668, 1.00623059, 2.64196896,
       2.77173231, 1.        , 1.95256242, 2.9291637 , 2.        ,
       2.93470612, 2.93470612, 2.        , 1.95256242, 1.        ,
       2.77173231, 1.00623059, 2.86818758, 2.66129668, 0.46097722,
       1.        , 2.83310783, 2.91247318, 1.18848643, 2.        ,
       2.14767316, 2.        , 2.9291637 , 2.83310783, 1.        ,
       2.64196896, 2.86818758, 2.71661554, 1.        , 2.        ])

Finally I can create an array of shape (3,16) by executing

distances = np.zeros((n_atoms, n_max_nbh), dtype=np.float32)
distances[mask] = dist
distances
array([[2.1476731 , 2.        , 1.1884865 , 2.9124732 , 1.        ,
        0.46097723, 2.7166154 , 2.6612966 , 1.0062306 , 2.641969  ,
        2.7717323 , 1.        , 1.9525625 , 2.9291637 , 2.        ,
        2.9347062 ],
       [2.9347062 , 2.        , 1.9525625 , 1.        , 2.7717323 ,
        1.0062306 , 2.8681877 , 2.6612966 , 0.46097723, 1.        ,
        2.833108  , 2.9124732 , 1.1884865 , 2.        , 2.1476731 ,
        0.        ],
       [2.        , 2.9291637 , 2.833108  , 1.        , 2.641969  ,
        2.8681877 , 2.7166154 , 1.        , 2.        , 0.        ,
        0.        , 0.        , 0.        , 0.        , 0.        ,
        0.        ]], dtype=float32)
distances.shape
(3,16)
4

Thanks for the code snippet. I don’t see any numpy-specific operation which wouldn’t work in PyTorch, too. Based on your previous post:

it sounded as if the approach works fine in numpy but is not supported in PyTorch. Is this indeed the case (and I’m missing something)?