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.
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.]])
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)
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.]])
arr.flatattribute.