Replacing part of a 2D numpy array using indexing

Question

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 elements row i and j and column 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 array S with another array of same dimensions (tmp_arr) python does not give an error but it also does not do anything meaning that the elements of S remain unchanged and no error message is displayed.

>>> tmp_arr
Out[3]: 
array([[ 555.,  0.],
       [ 0.,  555.]])

>>> S[:, [i, j]][[i, j], :] = tmp_arr

and what I get is the same matrix:

>>> S
Out[4]: 
array([[ 1.,  0.,  0.],
      [ 0.,  3.,  0.],
      [ 0.,  0.,  9.]])

Obviously the following would work but I am 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 experiences.

Answer 1

You could use np.ix_ to construct the desired index arrays:

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]])

---

Using np.ix_ is good for making assignments to S, but note there are faster ways to select the 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, S[np.ix_(...)] = ... does not use chained indexing, so S.__setitem__ gets called and the assignment affects S. In contrast, S[:, [i, j]] returns a copy of a subarray of S, so assigning to S[:, [i, j]][[i, j], :] affects only this copy of the subarray, not S itself. Since no reference to this copy of the subarray is maintained, Python throws away the copy after the assignment is made, so the assignment is lost. That is why chained indexing is no good for assigning to S.