Numpy: Assignment and Indexing as Matlab

Question

Sometimes is useful to assign arrays with one index only. In Matlab this is straightforward:

M = zeros(4);
M(1:5:end) = 1
M =

   1   0   0   0
   0   1   0   0
   0   0   1   0
   0   0   0   1

Is there a way to do this in Numpy? First I thought to flatten the array, but that operation doesn't preserve the reference, as it makes a copy. I tried with ix_ but I couldn't manage to do it with a relatively simple syntax.

Answer 1

You could try numpy.ndarray.flat, which represents an iterator that you can use for reading and writing into the array.

>>> M = zeros((4,4))
>>> M.flat[::5] = 1
>>> print(M)
array([[ 1.,  0.,  0.,  0.],
       [ 0.,  1.,  0.,  0.],
       [ 0.,  0.,  1.,  0.],
       [ 0.,  0.,  0.,  1.]])

Note that in numpy the slicing syntax is [start:stop_exclusive:step], as opposed to Matlab's (start:step:stop_inclusive).

Based on sebergs comment it might be important to point out that Matlab stores matrices in column major, while numpy arrays are row major by default.

>>> M = zeros((4,4))
>>> M.flat[:4] = 1
>>> print(M)
array([[ 1.,  1.,  1.,  1.],
       [ 0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.]])

To get Matlab-like indexing on the flattened array you will need to flatten the transposed array:

>>> M = zeros((4,4))
>>> M.T.flat[:4] = 1
>>> print(M)
array([[ 1.,  0.,  0.,  0.],
       [ 1.,  0.,  0.,  0.],
       [ 1.,  0.,  0.,  0.],
       [ 1.,  0.,  0.,  0.]])

Answer 2

You could do this using list indices:

M = np.zeros((4,4))
M[range(4), range(4)] = 1
print M
# [[ 1.  0.  0.  0.]
#  [ 0.  1.  0.  0.]
#  [ 0.  0.  1.  0.]
#  [ 0.  0.  0.  1.]]

In this case you could also use np.identity(4)

Answer 3

Another way using unravel_index

>>> M = zeros((4,4));
>>> M[unravel_index(arange(0,4*4,5),(4,4))]= 1
>>> M
array([[ 1.,  0.,  0.,  0.],
       [ 0.,  1.,  0.,  0.],
       [ 0.,  0.,  1.,  0.],
       [ 0.,  0.,  0.,  1.]])