Scatter-Like Operation in Numpy

Question

I am faced with the following problem:

I have a numpy array A of shape (*S, N, N) where S is an arbitrary tuple and N is some positive integer. Another array I of shape (*S, 2) represents indices in A. The values in I are integers in {0, ..., N-1}.

I would like to set

A[i1, ...., ik, I[i1, ..., ik, 0], I[i1, ...., ik, 1]] := 0 (Pseudocode)

for all valid indices i1, ..., ik.

For instance, consider the following example where N=3 and S=(2):

A = np.array([[[ 1,  2,  3],
               [ 4,  5,  6],
               [ 7,  8,  9]],

              [[10, 11, 12],
               [13, 14, 15],
               [16, 17, 18]]])

I = np.array([[0, 2],
              [2, 1]])

The desired output in this case would be

np.array([[[ 1,  2,  0],
           [ 4,  5,  6],
           [ 7,  8,  9]],

          [[10, 11, 12],
           [13, 14, 15],
           [16, 0, 18]]])

Note that the 3 at index (0, 0, 2) (corresponding to the value I[0] being np.array([0, 2])) and the 17 at index (1, 2, 1) (corresponding to the value I[1] being np.array([2, 1])) were changed to 0.

This looks a little like the scatter operation in PyTorch. However, it doesn't exactly match that problem. Additionally, numpy doesn't provide a scatter method. I have worked through numpy's documentation on indexing and a variety of methods to distribute values. Nonetheless, I haven't managed to come up with a clever approach to solve this problem (i.e. one without loops).

I would be grateful for suggestions!

user_13091658 · Answer

Interesting problem!

I believe you can simply reshape A and I to (-1,N,N) and (-1,2) respectively and always handle the problem in (K x N x N), then simply reshape back to to original dimension. My proposed solution is as follows:

import numpy as np

# Just extend the problem a little because why not
A = np.array([[[ 1,  2,  3],
               [ 4,  5,  6],
               [ 7,  8,  9]],

              [[10, 11, 12],
               [13, 14, 15],
               [16, 17, 18]],
            
              [[19, 20, 21],
               [22, 23, 24],
               [25, 26, 27]],
])

I = np.array([[0, 2],
              [2, 1],
              [1, 1]])


B = np.stack([A,A])
C = np.array([
    [[0,2],
     [2,1],
     [1,1]],
    
    [[0,0],
     [1,1],
     [2,2]],
])

def scatterlike(A: np.ndarray, I: np.ndarray, target: float=0):
    A_ = A.copy().reshape(-1, *A.shape[-2:])
    A_[(range(len(A_)),*I.reshape(-1,2).T.tolist())] = target
    return A_.reshape(A.shape)

I did a A.copy() in the code as it made it easier to debug, but you can do everything in place if you'd like. I am pretty sure this code works for (*S, N, M) as well, which is nice. I avoid indexing using np.ndarray types since the indexing behavior is different for np.ndarray and other iterables.

Here are some outputs:

# Simplest case  
scatterlike(A[0], I[0])
array([[1, 2, 0],
       [4, 5, 6],
       [7, 8, 9]])

# Extended version of the example you (OP) provided
scatterlike(A, I)
array([[[ 1,  2,  0],
        [ 4,  5,  6],
        [ 7,  8,  9]],

       [[10, 11, 12],
        [13, 14, 15],
        [16,  0, 18]],

       [[19, 20, 21],
        [22,  0, 24],
        [25, 26, 27]]])

# Higher dimensions, B.shape is (2,3,3,3)
scatterlike(B, C)
array([[[[ 1,  2,  0],
         [ 4,  5,  6],
         [ 7,  8,  9]],

        [[10, 11, 12],
         [13, 14, 15],
         [16,  0, 18]],

        [[19, 20, 21],
         [22,  0, 24],
         [25, 26, 27]]],


       [[[ 0,  2,  3],
         [ 4,  5,  6],
         [ 7,  8,  9]],

        [[10, 11, 12],
         [13,  0, 15],
         [16, 17, 18]],

        [[19, 20, 21],
         [22, 23, 24],
         [25, 26,  0]]]])

# Even higher dimensions, B[None].repeat(2,0) is of shape (2,2,3,3,3)
# Output should be like above but repeated
scatterlike(B[None].repeat(2,0), C[None].repeat(2,0))
array([[[[[ 1,  2,  0],
          [ 4,  5,  6],
          [ 7,  8,  9]],

         [[10, 11, 12],
          [13, 14, 15],
          [16,  0, 18]],

         [[19, 20, 21],
          [22,  0, 24],
          [25, 26, 27]]],


        [[[ 0,  2,  3],
          [ 4,  5,  6],
          [ 7,  8,  9]],

         [[10, 11, 12],
          [13,  0, 15],
          [16, 17, 18]],

         [[19, 20, 21],
          [22, 23, 24],
          [25, 26,  0]]]],



       [[[[ 1,  2,  0],
          [ 4,  5,  6],
          [ 7,  8,  9]],

         [[10, 11, 12],
          [13, 14, 15],
          [16,  0, 18]],

         [[19, 20, 21],
          [22,  0, 24],
          [25, 26, 27]]],


        [[[ 0,  2,  3],
          [ 4,  5,  6],
          [ 7,  8,  9]],

         [[10, 11, 12],
          [13,  0, 15],
          [16, 17, 18]],

         [[19, 20, 21],
          [22, 23, 24],
          [25, 26,  0]]]]])

Bonus

Just for fun, I think you can even turn the function to work for (*S,*G) (by getting the shape information from I) as well, that is, instead of only 2D arrays of shape (N x N), you can have like (N x M x L x ...). Also you could specify what each target value should be:

from numpy.typing import ArrayLike
from typing import Union
def scatterlike_general(A: np.ndarray, I: np.ndarray, target: Union[ArrayLike,float] = 0):
    A_ = A.reshape(-1, *A.shape[-I.shape[-1]:])
    A_[(range(len(A_)),*I.reshape(-1,I.shape[-1]).T.tolist())] = target
    return A_.reshape(A.shape)

It should be noted that I have not tested scatterlike_general, but I think it should work.

user_15306606 · Answer

You can access the dimensions of A represented by S using a flattened mesh grid. To access the appropriate locations in I, you can use the same flattened mesh grid. He's some code that implements this idea.

Here's an additional more complex example (sorry for not including the whole array A at the output):

Scatter-Like Operation in Numpy

2 Answers

Bonus