CODEWITHSUNDEEP
pythonbash
B_runo
B_runo

Reputation: 159

Alternative to bash (awk command) with python

Context : I run calculations on a program that gives me result files. On these result files (extension .h5), I can apply a python code (I cannot change this python code) such that it gives me a square matrix :

oneptdm.py resultfile.h5

gives me for example :

1   2  3  4
5   6  7  8
9  10 11 12
13 14 15 16
points groups
1
2
3
...

in a file called oneptdm.dat

I want to grep the diagonal of this matrix. Usually I use simply bash:

awk '{ for (i=0; i<=NF; i++) if (NR >= 1 && NR == i) print i,$(i) }' oneptdm.dat > diagonal.dat

But for x reason, I have to do it with python now. How can I do that ? I can of course use "subprocess" to use awk again but I would like to know if there is an alternative way to do that with a python script, version 2.6. The result should be : 

(line)  (diagonal element)
1       1
2       6
3       11
4       16

Upvotes: 0

Views: 2459

Answers (2)

mhawke
mhawke

Reputation: 87084

This should do the trick. It does assume that the file begins with a square matrix, and that assumption is used to limit the number of lines read from the file.

with open('oneptdm.dat') as f:
    line = next(f).split()
    for i in range(len(line)):
        print('{0}\t{1}'.format(i+1, line[i]))
        try:
            line = next(f).split()
        except StopIteration:
            break

Output for your sample file:

1   1
2   6
3   11
4   16

Upvotes: 1

Aleksei Shestakov
Aleksei Shestakov

Reputation: 2538

You can try something like this:

with open('oneptdm.dat') as f:
    for i, l in enumerate(f):
        print '%d\t%s' % (i + 1, l.split()[i])

Upvotes: 2

Related Questions