I have the following arrays.
b0 = np.zeros((2, 3, 4))
array([[[0, 0, 0, 0],
[0, 0, 0, 0],
[0, 0, 0, 0]],
[[0, 0, 0, 0],
[0, 0, 0, 0],
[0, 0, 0, 0]]])
b1 = np.array([
[1, 0, 2],
[0, 0, 1]
])
b2 = np.array([
[100, 0, 0],
[200, 300, 0]
])
My goal is to update values of b0 using indices specified in b1 and new values given by b2. Specifically, I want to make these updates
b0[0, 0, b1[0,0]] = b2[0,0]
b0[0, 1, b1[0,1]] = b2[0,1]
b0[0, 2, b1[0,2]] = b2[0,2]
b0[1, 0, b1[1,0]] = b2[1,0]
b0[1, 1, b1[1,1]] = b2[1,1]
b0[1, 2, b1[1,2]] = b2[1,2]
b0
array([[[ 0, 100, 0, 0],
[ 0, 0, 0, 0],
[ 0, 0, 0, 0]],
[[200, 0, 0, 0],
[300, 0, 0, 0],
[ 0, 0, 0, 0]]])
How do I do this dynamically?
In [51]: b1 = np.array([
...: [1, 0, 2],
...: [0, 0, 1]
...: ])
...:
...: b2 = np.array([
...: [100, 0, 0],
...: [200, 300, 0]
...: ])
In [52]: b0 = np.zeros((2,3,4),int)
In [53]: b0[np.arange(2)[:,None], np.arange(3), b1] = b2
In [54]: b0
Out[54]:
array([[[ 0, 100, 0, 0],
[ 0, 0, 0, 0],
[ 0, 0, 0, 0]],
[[200, 0, 0, 0],
[300, 0, 0, 0],
[ 0, 0, 0, 0]]])
np.arange(2)[:,None], np.arange(3) broadcast together to index a (2,3) block, same as b1 shape.
meshgrid lets us compare my indexing and yours:
In [58]: np.meshgrid(np.arange(2), np.arange(3), sparse=True, indexing='ij')
Out[58]:
[array([[0],
[1]]),
array([[0, 1, 2]])]
In [59]: np.meshgrid(np.arange(2), np.arange(3), sparse=False, indexing='ij')
Out[59]:
[array([[0, 0, 0],
[1, 1, 1]]),
array([[0, 1, 2],
[0, 1, 2]])]
Just realized I can do this
n = b0.shape[0]
k = b2.shape[1]
i1 = np.repeat(np.arange(n), k)
i2 = np.tile(np.arange(k), n)
b0[i1, i2, b1[i1, i2]] = b2[i1, i2]
but there's probably a better way to do this.
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With