How to assign different values to elements at different positions in one array based on multiple conditions

[ijku]

Hi everyone,

As the title suggests, I want to assign different values to elements at different positions in an array based on multiple conditions simultaneously. Here’s a simplified example: given an array
m = [ [0, 1, 2, 3, 4], [5, 6, 7, 8, 9], [10, 11, 12, 13, 14] ]
I want to modify the elements adjacent to the positions where m == 2, 6, 12 by setting them to 2, 6, and 12, respectively. The result should be:
m = [ [0, 2, 2, 2, 4], [6, 6, 6, 8, 9], [10, 12, 12, 12, 14] ]

I used vmap and where functions to achieve this, with code as follows:

import jax
import jax.numpy as np

def update_by_multi_conditions(arr, indexs):
    def get_partial_update(arr, index):
        mask = (np.abs(arr-index)<=1)
        upd = np.where(mask, index, 0)
        return upd, mask

    vmap_update_fn = jax.vmap(get_partial_update, in_axes=[None, 0])

    upds, masks = vmap_update_fn(arr, indexs)
    updated_mask = np.any(masks, axis=0)
    updated_results = upds.sum(axis=0) 
    result = np.where(updated_mask, updated_results, arr)

    return result

m = np.arange(15).reshape(3, 5)
n = np.array([2,6,12])
l = update_by_multi_conditions(m, n)
print(l)

My general approach is to use vmap to locate all positions, modify them individually, and assign a value of 0 to positions that haven’t been modified. Then, I sum everything. Finally, I use the where function to keep original values at unmodified positions, while for modified positions, I assign the accumulated values from the previous step.

However, I feel this method is overly complicated, and there might be a more efficient solution. Does anyone have any suggestions?

[jakevdp]

For this basic scenario, you could do something based on jax.scipy.signal.convolve:

import jax.numpy as jnp
import jax.scipy as jsp

m = jnp.array([[0, 1, 2, 3, 4],
               [5, 6, 7, 8, 9],
               [10, 11, 12, 13, 14]])
vals = jnp.array([2, 6, 12])

masked = jnp.where(jnp.isin(m, vals), m, 0)
convolved = jsp.signal.convolve(masked, jnp.ones((1, 3)), mode='same').astype(m.dtype)
result = jnp.where(convolved != 0, convolved, m)
print(result)

[[ 0  2  2  2  4]
 [ 6  6  6  8  9]
 [10 12 12 12 14]]

This would have to be modified if 0 is one of the values you want to consider, and you'd also have to be more clear about what you expect in the case of overlaps, but I think this is roughly the right approach.

[ijku]

Jake, thank you for your response; it’s really insightful! However, the situation I’m dealing with is actually more complex, and the example I gave may have oversimplified :) Sorry about that. Using the convolve , the shape of the modified positions in the original array would align with the convolve’s right-hand input, but in my case, it needs to be more flexible.

The problem can be roughly re-described like this:
in array a find the positions of the highlighted (red) elements (these highlighted positions form a connected region).

a = jnp.array([100,100,0,0,0,0,100,100,100,0,0,0,5,100,2,2,2,0,5,5,101,101,101,3,5,5,101,101,101,3,5,5,101,101,101,3])
a = jnp.reshape(a, [6,6])

Then, in another array b of the same shape, modify the values in the corresponding region based on the center of that region.

b = jnp.array([12,48,23,2,45,45,99,0,1,2,3,4,5,7,7,8,93,57,6,91,33,76,10,9,7,25,38,1,33,55,8,89,21,81,71,66])
b = jnp.reshape(b, [6,6])

After modification, b should look like this:

I just made up some random data. Do you have any good ideas or suggestions for this problem? Thanks for your time and help.