Replacing part of a 2D numpy array using indexing
I am trying to replace part of a 2D numpy array named "S" as a function of i and j. Given S as:
>>> S
Out[1]:
array([[ 1., 0., 0.],
[ 0., 3., 0.],
[ 0., 0., 9.]]
for i = 0 and j = 1, I can access the elements of the rows i and j and the columns i and j using the following syntax:
>>> S[:, [i, j]][[i, j], :]
Out[2]:
array([[ 1., 0.],
[ 0., 3.]])
Now, when I try to replace the same elements of the S array with another array of the same size (tmp_arr), python does not give an error, but it also does nothing, which means the S elements remain the same and no error message is displayed.
>>> tmp_arr
Out[3]:
array([[ 555., 0.],
[ 0., 555.]])
>>> S[:, [i, j]][[i, j], :] = tmp_arr
and I get the same matrix:
>>> S
Out[4]:
array([[ 1., 0., 0.],
[ 0., 3., 0.],
[ 0., 0., 9.]])
Obviously the following will work, but I'm looking for an elegant solution:
S[i, i] = tmp_arr[0, 0] S[i, j] = tmp_arr[0, 1] S[j, i] = tmp_arr[1, 0] S[j, j] = tmp_arr[1, 1]
I appreciate your comments and impressions.
source to share
You can use np.ix_
to build the required arrays of indices:
In [91]: S[np.ix_([i,j],[i,j])]
Out[91]:
array([[1, 0],
[0, 3]])
In [92]: tmp_arr = np.eye(2)*555
In [93]: tmp_arr
Out[93]:
array([[ 555., 0.],
[ 0., 555.]])
In [94]: S[np.ix_([i,j],[i,j])] = tmp_arr
In [95]: S
Out[95]:
array([[555, 0, 0],
[ 0, 555, 0],
[ 0, 0, 9]])
Usage np.ix_
is useful for assignments S
, but note that there are faster ways to select a subarray:
In [99]: %timeit S.take([i, j], axis=1).take([i, j], axis=0)
100000 loops, best of 3: 3.32 µs per loop
In [97]: %timeit S[:, [i, j]][[i, j], :]
100000 loops, best of 3: 8.8 µs per loop
In [96]: %timeit S[np.ix_([i,j],[i,j])]
100000 loops, best of 3: 13 µs per loop
But unlike these other methods, it S[np.ix_(...)] = ...
does not use the indexing chain, so it S.__setitem__
is called, and the assignment affects S
. On the contrary, it S[:, [i, j]]
returns a copy of the subarray, S
so the assignment S[:, [i, j]][[i, j], :]
only affects that copy of the subarray, not S
. Since a reference to this copy of the subarray is not maintained, Python discards the copy after the job is done, so the assignment is lost. This is why the index chain is not suitable for assignment S
.
source to share