Sum over selected numpy.ndarray column and write to a file I am trying to write the sum of selected NumPy array column and write to a file.
The following code snippet is working fine:

for i in range(int(ldata[2])):
    s_u.append(dos2[i][1])
    s_d.append(dos2[i][2])
    p_u.append(np.sum(dos2[i,[3,5,7]]))
    p_d.append(np.sum(dos2[i,[4,6,8]]))
    d_u.append(np.sum(dos2[i,[9,11,13,15,17]]))
    d_d.append(np.sum(dos2[i,[10,12,14,16,18]]))
with open("dos.dat","w") as ados:
    for i in range(int(ldata[2])):
        ados.write(str(atom[0,i,0]-float(ldata[3]))+" ")
        ados.write(str(s_u[i])+" ")
        ados.write(str(s_d[i])+" ")
        ados.write(str(p_u[i])+" ")
        ados.write(str(p_d[i])+" ")
        ados.write(str(d_u[i])+" ")
        ados.write(str(d_d[i])+" ")
        ados.write("\n")
print(type(dos2))   #&lt;class 'numpy.ndarray'&gt;
print(type(s_u))    #&lt;class 'list'&gt;


But I am feeling like I am overdoing things. Also, I am not checking how many columns (j) there are in dos2[i,[j]] (which may vary).

Can I get some help to make it better?
I will assume that all of the ?_? lists are empty when you start.  Here are my suggestions:


Rather than using lists, I would just make the ?_? variables numpy arrays.  You can slice the values you want from dos2 using syntax like s_u = dos2[:int(ldata[2]), [3,5,7]].
Your indexing is really slices with a step size of two.  So [3, 5, 7] is 3:9:2, [4,6,8] is 4:10:2, [9,11,13,15,17] is 9:19:2, and [10,12,14,16,18] is 10:20:2.  Numpy is much faster with these sorts of slices rather than using the direct indexing you are using since with these slices numpy can avoid making a copy.
Since everything is done in steps of 2, I think it would make it even easier to first reshape dos2 to a 3D array, where the last dimension determines whether it is a ?_u or ?_d variables.  Then you can use unpacking to put the columns in the right place.
You can use np.column_stack to combine all of your 1D arrays into one big 2D array.  This can then be written in one step using np.savetxt.
Better yet, you can use a slice for the last index of atom to end up with all 2D arrays, then just hstack them to get one big array.  This avoids having to unpack at all.
I prefer the arr.sum() syntax of np.sum(arr) since I think it makes the flow easier to follow for people who read left-to-right.


So here is how I would implement what you are doing (with comments and dummy data):

# Dummy test data
dos2 = np.random.random((1000, 19))
atom = np.random.random((2, 1000, 2))
ldata = [None, None, '100', '1.1']

# We re-use this value several times, so put it in a variable
maxind = int(ldata[2])

# Extract the part of dos2 we want, and reshape it so the last index
# determines whether it is a ?_u (index==0) or ?_d (index==1) value
dos2_r = dos2[:maxind, 1:19].reshape(-1, 9, 2)

# Extract the part of atom we want
# We use slicing instead of indexing for the last index to keep it 2D
atom_p = atom[0,:maxind,:1]-float(ldata[3])

# Slice out the parts of dos2 we want
s_ud = dos2_r[:, 0, :]
p_ud = dos2_r[:, 1:4, :].sum(axis=1)
d_ud = dos2_r[:, 4:, :].sum(axis=1)

# Combine the arrays into one big array
cols = np.hstack([atom_p, s_ud, p_ud, d_ud])

# Write to a file
np.savetxt('dos.dat', cols, delimiter=' ')


Edit:

OP asked my to explain how the slicing is working.  I think the easiest way to explain is to walk through a toy example.  Lets start with a 2D array. That I think makes it easy to follow what is going on:

&gt;&gt;&gt; dos2 = np.vstack([np.arange(12)+100*i for i in range(20)]).T
&gt;&gt;&gt; print(dos2)
[[   0  100  200  300  400  500  600  700  800  900 1000 1100 1200 1300 1400 1500 1600 1700 1800 1900]
 [   1  101  201  301  401  501  601  701  801  901 1001 1101 1201 1301 1401 1501 1601 1701 1801 1901]
 [   2  102  202  302  402  502  602  702  802  902 1002 1102 1202 1302 1402 1502 1602 1702 1802 1902]
 [   3  103  203  303  403  503  603  703  803  903 1003 1103 1203 1303 1403 1503 1603 1703 1803 1903]
 [   4  104  204  304  404  504  604  704  804  904 1004 1104 1204 1304 1404 1504 1604 1704 1804 1904]
 [   5  105  205  305  405  505  605  705  805  905 1005 1105 1205 1305 1405 1505 1605 1705 1805 1905]
 [   6  106  206  306  406  506  606  706  806  906 1006 1106 1206 1306 1406 1506 1606 1706 1806 1906]
 [   7  107  207  307  407  507  607  707  807  907 1007 1107 1207 1307 1407 1507 1607 1707 1807 1907]
 [   8  108  208  308  408  508  608  708  808  908 1008 1108 1208 1308 1408 1508 1608 1708 1808 1908]
 [   9  109  209  309  409  509  609  709  809  909 1009 1109 1209 1309 1409 1509 1609 1709 1809 1909]
 [  10  110  210  310  410  510  610  710  810  910 1010 1110 1210 1310 1410 1510 1610 1710 1810 1910]
 [  11  111  211  311  411  511  611  711  811  911 1011 1111 1211 1311 1411 1511 1611 1711 1811 1911]]


This has 12 values in the first dimension and 20 in the second, so it has some extra data rows and columns.  I chose this particular array since, as you can see, for value xyy or xxyy, the x or xx value is the column number and the yy value is the row number.  This makes it easy to keep track of what the manipulations are doing.

Let's make maxind = 10 for simplicity.  So the first thing I do is slice out the data you don't want.  You never use the first column, and you never use past the 18th column.  Also, you never use past maxind rows.  So I throw out everything besides that.  I will call the result dos2_p.  Note that with slicing, for arr[0:n] the 0 is implied so it can be written as arr[:n], and for arr[n:len(arr)] the len(arr) is implied so it can be written as arr[n:].  arr[:] means "everthing", and arr[:, n] means "everything from dimension 0 and the nth index of dimension 1".  arr[..., n] means "the nth index of the last dimension and everything from the rest of the dimensions, however many that might be"

&gt;&gt;&gt; maxind = 10
&gt;&gt;&gt; dos2_p = dos2[:maxind, 1:19]
&gt;&gt;&gt; print(dos2_p)
[[ 100  200  300  400  500  600  700  800  900 1000 1100 1200 1300 1400 1500 1600 1700 1800]
 [ 101  201  301  401  501  601  701  801  901 1001 1101 1201 1301 1401 1501 1601 1701 1801]
 [ 102  202  302  402  502  602  702  802  902 1002 1102 1202 1302 1402 1502 1602 1702 1802]
 [ 103  203  303  403  503  603  703  803  903 1003 1103 1203 1303 1403 1503 1603 1703 1803]
 [ 104  204  304  404  504  604  704  804  904 1004 1104 1204 1304 1404 1504 1604 1704 1804]
 [ 105  205  305  405  505  605  705  805  905 1005 1105 1205 1305 1405 1505 1605 1705 1805]
 [ 106  206  306  406  506  606  706  806  906 1006 1106 1206 1306 1406 1506 1606 1706 1806]
 [ 107  207  307  407  507  607  707  807  907 1007 1107 1207 1307 1407 1507 1607 1707 1807]
 [ 108  208  308  408  508  608  708  808  908 1008 1108 1208 1308 1408 1508 1608 1708 1808]
 [ 109  209  309  409  509  609  709  809  909 1009 1109 1209 1309 1409 1509 1609 1709 1809]]


As you can see, each rows start with the 1?? value and end with the 18?? value, which means it starts with columns 1 and ends with column 18, and the maximum ?yy value is 09, which means it ends with row 9.

Now I want to make it so the even rows are in one dimension, and the odd rows are in another.  I can do that with reshape.  It can take some trial-and-error to get the reshape to work the way you want.  I don't actually have a good conceptual way to do this, I just create a dummy example like this and try different things until I get the layout I want.  In this case, we want to set the last dimension to have a shape of 2, the first to have the same shape as currently (maxind), and we can let numpy figure out the second dimension (although we know it should be 9).

&gt;&gt;&gt; dos2_r = dos2_p.reshape(maxind, -1, 2)
&gt;&gt;&gt; print(dos2_r[..., 0])
[[ 100  300  500  700  900 1100 1300 1500 1700]
 [ 101  301  501  701  901 1101 1301 1501 1701]
 [ 102  302  502  702  902 1102 1302 1502 1702]
 [ 103  303  503  703  903 1103 1303 1503 1703]
 [ 104  304  504  704  904 1104 1304 1504 1704]
 [ 105  305  505  705  905 1105 1305 1505 1705]
 [ 106  306  506  706  906 1106 1306 1506 1706]
 [ 107  307  507  707  907 1107 1307 1507 1707]
 [ 108  308  508  708  908 1108 1308 1508 1708]
 [ 109  309  509  709  909 1109 1309 1509 1709]]
&gt;&gt;&gt; print(dos2_r[..., 1])
[[ 200  400  600  800 1000 1200 1400 1600 1800]
 [ 201  401  601  801 1001 1201 1401 1601 1801]
 [ 202  402  602  802 1002 1202 1402 1602 1802]
 [ 203  403  603  803 1003 1203 1403 1603 1803]
 [ 204  404  604  804 1004 1204 1404 1604 1804]
 [ 205  405  605  805 1005 1205 1405 1605 1805]
 [ 206  406  606  806 1006 1206 1406 1606 1806]
 [ 207  407  607  807 1007 1207 1407 1607 1807]
 [ 208  408  608  808 1008 1208 1408 1608 1808]
 [ 209  409  609  809 1009 1209 1409 1609 1809]]


See that index 0 of the last dimension has all the odd columns, while dimension 1 has all the even columns.

Now let's do a min in this case instead of a sum.  The semantics are the same, but it is easier to follow in my opinion.  We will take the min of dimension 1, which is what used to be columns, but now is the columns split into even and odd:

&gt;&gt;&gt; dos2_m = dos2_r.min(axis=1)
&gt;&gt;&gt; print(dos2_m)
[[100 200]
 [101 201]
 [102 202]
 [103 203]
 [104 204]
 [105 205]
 [106 206]
 [107 207]
 [108 208]
 [109 209]]


As you can see, it preserves the row number, and it preserves whether the column is even or odd, but it collapses across even columns and odd columns.

Now let's split it the way it is split in your code:

&gt;&gt;&gt; ms_ud = dos2_r[:, 0, :]
&gt;&gt;&gt; mp_ud = dos2_r[:, 1:4, :].min(axis=1)
&gt;&gt;&gt; md_ud = dos2_r[:, 4:, :].min(axis=1)
&gt;&gt;&gt; print(ms_ud)
[[100 200]
 [101 201]
 [102 202]
 [103 203]
 [104 204]
 [105 205]
 [106 206]
 [107 207]
 [108 208]
 [109 209]]
&gt;&gt;&gt; print(mp_ud)
[[300 400]
 [301 401]
 [302 402]
 [303 403]
 [304 404]
 [305 405]
 [306 406]
 [307 407]
 [308 408]
 [309 409]]
&gt;&gt;&gt; print(md_ud)
[[ 900 1000]
 [ 901 1001]
 [ 902 1002]
 [ 903 1003]
 [ 904 1004]
 [ 905 1005]
 [ 906 1006]
 [ 907 1007]
 [ 908 1008]
 [ 909 1009]]


So you can see it takes the same columns you did in your original, but keeps them in a single 2D array.  And it greatly simplifies the code and speeds up the computation as well.