Iterating through N lines in a for loop before moving on to next loop

Question

I posted a question on How to interate over N lines of a text file in Python but I wanted to make the program more applicable to N lines (e.g. instead of having to collect two line at a time I wanted to take four or six lines at a time.

I want to turn this disp.txt :

116  C             0.12 -0.91  0.39    -0.40  0.31  0.85    -0.66 -0.18 -0.22
117  O             0.00 -0.02  0.00    -0.05  0.05  0.12    -0.57 -0.26 -0.29
116  C            -0.03 -0.04  0.00     0.01  0.09  0.19    -0.71 -0.21 -0.26
117  O            -0.14  0.88 -0.45     0.47 -0.33 -0.79     0.57  0.16  0.19

To this:

vibration 1
0.12  -0.91  0.39
0.0  -0.02  0.0
vibration 2
-0.4  0.31  0.85
-0.05  0.05  0.12
vibration 3
-0.66  -0.18  -0.22
-0.57  -0.26  -0.29
vibration 4 
-0.03 -0.04 0.00 
-0.14 0.88 -0.45
vibration 5 ...

Below is what I have so far

vibration = 1
with open('disp1.txt','r') as f, open('disp2.txt', 'w') as out:

 while True:
  next_lines = list(islice(f,2))
  m = next_lines
  for i in m:
   s = i.strip().split('\n')
   for j in s:
    p = j.split()
    v = len(p)
    for u in range(2,5,3):
     out.write('vibration {}\n'.format(vibration))
     out.write(' '.join(p[u:u+3])+'\n')
     vibration += 1
    for x in range(5,8,3):
     out.write(' '.join(p[x:x+3])+'\n')
    for y in range(8,11,3):
      out.write(' '.join(p[y:y+3])+'\n')
  if not next_lines:
   break

Which is showing this:

vibration 1
0.12 -0.91 0.39
-0.40 0.31 0.85
-0.66 -0.18 -0.22
vibration 2
0.00 -0.02 0.00
-0.05 0.05 0.12
-0.57 -0.26 -0.29
vibration 3
-0.03 -0.04 0.00
0.01 0.09 0.19
-0.71 -0.21 -0.26
vibration 4
-0.14 0.88 -0.45
0.47 -0.33 -0.79
0.57 0.16 0.19

So it's reading across the page instead of reading the first two lines of first three columns of numbers and then proceeding to the next set of three and so on. By using the isplice it reads two lines at a time but for reasons I don't understand both lines are not passed through the first loop (for u in range(2,5,3):) before moving onto the next one (for x in range(5,8,3):) Instead each line is put through each of the loops. Could anyone help me with this?

Thanks.

Answer 1

This should do what you want.

def group(lst, size = 3):
    """
    returns a list of lists where the inner lists are at most
    of length size

    [1, 2, 3, 4, 5, 6] -> [[1, 2, 3], [4, 5, 6]]
    group(list, int) -> list<list>
    """
    return [lst[i:i+size] for i in range(0, len(lst), size)]

def get_vibrations(lines):
    for vibration in zip(*[group(line) for line in lines]):
        yield '\n'.join(' '.join(v) for v in vibration)

def main(n):
    with open('in.txt', 'r') as fin, open('out.txt', 'w') as fout:
        i = 0
        for line in fin:
            line = line.strip().split()[2:]
            vibrations = [line] + [fin.readline().strip().split()[2:] for _ in range(n - 1)]

            # write each vibration to the file
            for i, v in enumerate(get_vibrations(vibrations), start = i + 1):
                fout.write('vibration {}\n{}\n'.format(i, v))

if __name__ == '__main__':
    main(4) # the number of lines per vibration

# Input
116  C             0.12 -0.91  0.39    -0.40  0.31  0.85    -0.66 -0.18 -0.22
117  O             0.00 -0.02  0.00    -0.05  0.05  0.12    -0.57 -0.26 -0.29
116  C            -0.03 -0.04  0.00     0.01  0.09  0.19    -0.71 -0.21 -0.26
117  O            -0.14  0.88 -0.45     0.47 -0.33 -0.79     0.57  0.16  0.19
116  C             0.26 -0.40 -0.83     0.26  0.31  0.25     0.33  0.71 -0.51
117  O            -0.19 -0.39 -0.36    -0.40  0.87  0.49     0.11  0.26 -0.57
116  C            -0.35  0.27  0.79    -0.03 -0.71 -0.18     0.93 -0.89 -0.81
117  O            -0.84  0.99  0.38    -0.27 -0.33  0.55     0.31  0.99 -0.18
116  C             0.76 -0.77  0.71     0.19 -0.39  0.46     0.30 -0.15  0.71
117  O            -0.97  0.31  0.58    -0.42  0.16 -0.33    -0.39 -0.17 -0.54

# Output
vibration 1
0.12 -0.91 0.39
0.00 -0.02 0.00
-0.03 -0.04 0.00
-0.14 0.88 -0.45
vibration 2
-0.40 0.31 0.85
-0.05 0.05 0.12
0.01 0.09 0.19
0.47 -0.33 -0.79
vibration 3
-0.66 -0.18 -0.22
-0.57 -0.26 -0.29
-0.71 -0.21 -0.26
0.57 0.16 0.19
vibration 4
0.26 -0.40 -0.83
-0.19 -0.39 -0.36
-0.35 0.27 0.79
-0.84 0.99 0.38
vibration 5
0.26 0.31 0.25
-0.40 0.87 0.49
-0.03 -0.71 -0.18
-0.27 -0.33 0.55
vibration 6
0.33 0.71 -0.51
0.11 0.26 -0.57
0.93 -0.89 -0.81
0.31 0.99 -0.18

If you also want to allow for not enough lines then you can add

from itertools import zip_longest

And change to this

def get_vibrations(lines):
    for vibration in zip_longest(*[group(line) for line in lines], fillvalue = ['null']):
        yield '\n'.join(' '.join(v) for v in vibration)

# Outputs | added to above output
vibration 7
0.76 -0.77 0.71
-0.97 0.31 0.58
null
null
vibration 8
0.19 -0.39 0.46
-0.42 0.16 -0.33
null
null
vibration 9
0.30 -0.15 0.71
-0.39 -0.17 -0.54
null
null

Answer 2

info = '''0.12 -0.91  0.39    -0.40  0.31  0.85    -0.66 -0.18 -0.22    0.00 -0.02  0.00    -0.05  0.05  0.12    -0.57 -0.26 -0.29    -0.03 -0.04  0.00    0.01  0.09  0.19    -0.71 -0.21 -0.26    -0.14  0.88 -0.45    0.47 -0.33 -0.79     0.57  0.16  0.19'''

new_list_1 = info.split('    ')
vib_1 = new_list_1[0] + '\n' + new_list_1[3]
vib_1 = 'vibration 1' + '\n' + vib_1
print vib_1

new_list_2 = info.split('    ')
vib_2 = new_list_2[1] + '\n' + new_list_2[4]
vib_2 = 'vibration 2' + '\n' + vib_2
print vib_2

new_list_3 = info.split('    ')
vib_3 = new_list_3[2] + '\n' + new_list_3[5]
vib_3 = 'vibration 3' + '\n' + vib_3
print vib_3