calculations for different columns in a numpy array

Question

I have a 2D array with filled with some values (column 0) and zeros (rest of the columns). I would like to do pretty much the same as I do with MS excel but using numpy, meaning to put into the rest of the columns values from calculations based on the first column. Here it is a MWE:

import numpy as np

a = np.zeros(20, dtype=np.int8).reshape(4,5)

b = [1, 2, 3, 4]

b = np.array(b)

a[:, 0] = b

# don't change the first column
for column in a[:, 1:]:
    a[:, column] = column[0]+1

The expected output:

array([[1, 2, 3, 4, 5],
       [2, 3, 4, 5, 6],
       [3, 4, 5, 6, 7],
       [4, 5, 6, 7, 8]], dtype=int8)

The resulting output:

array([[1, 0, 0, 0, 0],
       [1, 0, 0, 0, 0],
       [1, 0, 0, 0, 0],
       [1, 0, 0, 0, 0]], dtype=int8)

Any help would be appreciated.

Answer 1

Looping is slow and there is no need to loop to produce the array that you want:

>>> a = np.ones(20, dtype=np.int8).reshape(4,5)
>>> a[:, 0] = b
>>> a
array([[1, 1, 1, 1, 1],
       [2, 1, 1, 1, 1],
       [3, 1, 1, 1, 1],
       [4, 1, 1, 1, 1]], dtype=int8)
>>> np.cumsum(a, axis=1)
array([[1, 2, 3, 4, 5],
       [2, 3, 4, 5, 6],
       [3, 4, 5, 6, 7],
       [4, 5, 6, 7, 8]])

What went wrong

Let's start, as in the question, with this array:

>>> a
array([[1, 0, 0, 0, 0],
       [2, 0, 0, 0, 0],
       [3, 0, 0, 0, 0],
       [4, 0, 0, 0, 0]], dtype=int8)

Now, using the code from the question, let's do the loop and see what column actually is:

>>> for column in a[:, 1:]:
...   print(column)
... 
[0 0 0 0]
[0 0 0 0]
[0 0 0 0]
[0 0 0 0]

As you can see, column is not the index of the column but the actual values in the column. Consequently, the following does not do what you would hope:

a[:, column] = column[0]+1

Another method

If we want to loop (so that we can do something more complex), here is another approach to generating the desired array:

>>> b = np.array([1, 2, 3, 4])
>>> np.column_stack([b+i for i in range(5)])
array([[1, 2, 3, 4, 5],
       [2, 3, 4, 5, 6],
       [3, 4, 5, 6, 7],
       [4, 5, 6, 7, 8]])