[FIXED] Write a block of 3 lines repeatedly with calculated data into a text file using Python

Issue

I want to write data with headers into a file. The first three lines are unique, and can be considered as a 'block' which are then repeated with increments in x and y (0.12, 1) respectively. The data in the file should look like:

#X  #Y  Xmin    Ymin    Z
1   1   0.0000  0.000   0.0062
1   2   0.0000  0.350   0.0156
1   3   0.0000  0.750   0.0191
1   4   0.0000  1.000   0.0062
1   5   0.0000  1.350   0.0156
1   6   0.0000  1.750   0.0191
1   7   0.0000  2.000   0.0062
1   8   0.0000  2.350   0.0156
1   9   0.0000  2.750   0.0191
2   1   0.1200  0.000   0.0062
2   2   0.1200  0.350   0.0156
2   3   0.1200  0.750   0.0191
2   4   0.1200  1.000   0.0062
2   5   0.1200  1.350   0.0156
2   6   0.1200  1.750   0.0191
2   7   0.1200  2.000   0.0062
2   8   0.1200  2.350   0.0156
2   9   0.1200  2.750   0.0191
3   1   0.2400  0.000   0.0062
3   2   0.2400  0.350   0.0156
3   3   0.2400  0.750   0.0191
3   4   0.2400  1.000   0.0062
3   5   0.2400  1.350   0.0156
3   6   0.2400  1.750   0.0191
3   7   0.2400  2.000   0.0062
3   8   0.2400  2.350   0.0156
3   9   0.2400  2.750   0.0191

I tried to make the first three lines as 3 lists and write the first two columns and headers by two nested for loops but failed to write the repeating 3 line block.

l1 = [0.0000, 0.000, 0.0062]
l2 = [0.0000, 0.350, 0.0156]
l3 = [0.0000, 0.750, 0.0191]
pitch_x = 0.12
pitch_y = 1

with open('dataprep_test.txt', 'w') as f:
    f.write('#x #y  Xmin    Ymin    Z   \n')
    for i in range(1,4,1):
        k =1
        for j in range (1,4,1):
            d_x = pitch_x*(i-1)
            d_y = pitch_y*(j-1)
            f.write('%d %d  %f  %f  %f  \n'%(i,k,(l1[0]+d_x),(l1[1]+d_y), l1[2]))
            f.write('%d %d  %f  %f  %f  \n'%(i,k+1,(l2[0]+d_x),(l2[1]+d_y), l2[2]))
            f.write('%d %d  %f  %f  %f  \n'%(i,k+2,(l3[0]+d_x),(l3[1]+d_y), l3[2]))
            k=k+3

Is there a smarter way to do it using the python built-in functions and structures and methods (lists, dictionaries etc.)?

Solution

I'd just refactor the data generation into a generator function. You can also easily accept an arbitrary number of vectors.

def generate_data(initial_vectors, pitch_x, pitch_y, i_count=4, j_count=4):
    for i in range(i_count):
        for j in range(j_count):
            d_x = pitch_x * i
            d_y = pitch_y * j
            for k, (x, y, z) in enumerate(initial_vectors, 1):
                yield (i + 1, k, (x + d_x), (y + d_y), z)


def main():
    l1 = [0.0000, 0.000, 0.0062]
    l2 = [0.0000, 0.350, 0.0156]
    l3 = [0.0000, 0.750, 0.0191]
    with open('dataprep_test.txt', 'w') as f:
        f.write('#x #y  Xmin    Ymin    Z   \n')
        for i, k, x, y, z in generate_data([l1, l2, l3], pitch_x=0.12, pitch_y=1):
            f.write(f'{i:d} {k:d} {x:f} {y:f} {z:f}\n')


if __name__ == '__main__':
    main()

Furthermore, if a future version of your project might want to use JSON files instead, you could just json.dumps(list(generate_data(...)), etc.

Answered By - AKX