Selectively parsing large data file

Question

Task:

-I have run a simulation of a molecule vibrating, taking a "snapshot" of where my atoms are every 0.1 fempto-second

-Now I need to go through the output and grab all of these different sets of Cartesian coordinates (50 total).

-The geometry is in a nice format, however it has extra stuff with the same formatting in between geometry segments (see example).

Example (This is one of 50 identical segments from the file I'm trying to parse, only there are 16 atoms not 4.)

  Time in trajectory (femtosec)    3.300000D+00
 Total energy (au)  -2.716055737D+02
 Total angular momentum (h-bar)   5.485831060D-14
 Coordinates (Bohr)
 I=    1 X=  -1.573316541788D+00 Y=  -3.143098097327D-01 Z=  -9.488852008364D-01
 I=    2 X=  -1.549056004901D+00 Y=  -3.758762443395D-01 Z=   1.621875214114D+00
 I=    3 X=   7.633881398143D-01 Y=  -3.853052819189D-01 Z=   3.487630749614D-01
 I=    4 X=   2.390273062744D+00 Y=   1.832192143047D+00 Z=   6.235281069720D-01


 Momentum (sqrt(amu)*Bohr/sec)
 I=    1 X=  -3.536362458214D+13 Y=   2.431139678255D+13 Z=   1.413087646815D+13
 I=    2 X=  -2.526353110947D+13 Y=  -2.551987667221D+12 Z=   7.250485757030D+12
 I=    3 X=  -1.190415840625D+13 Y=  -1.311816871612D+13 Z=  -4.638293368564D+12
 I=    4 X=   6.566472312459D+12 Y=  -2.248922363477D+13 Z=  -4.123675084717D+11

I would like it if there was some way to loop over the file and read the coordinates, then skip the momentum information, and strip away all of the extra information as well such as I=, X= etc.

I was thinking it would be great if I could this into a 3D numpy array where my coordinates would be indexed by: (iteration number, atom number, and x y or z coordinate). I don't know if Regex would be much help here as the lines in momentum sections have identical formatting. Would any of you fine people be able to offer me a nudge in the right direction here?

Resources to go with an answer would be greatly appreciated I don't have a timeframe to work within so I am hoping to learn something here.

Answer 1

You can use numpy.genfromtxt() which accepts a generator and therefore allows some customization in the input you give. For example, here I am replacing D by e and ignoring all the lines starting with 'M' and 'C'. Additionally, I am using only columns (3, 5, 7), which contain the desired data. Finally, the array is reshaped to keep each "Momentum" separated from each other and accessible through the first index (assuming they will have 4 lines always):

gen = (r.replace('D','e') for r in open('test.txt') if not r[0] in ['C','M'])
a = np.genfromtxt(gen, skip_header=3, usecols=(3, 5, 7))
a = a.reshape(-1, 4, 3)

#[[[ -1.57331654e+00  -3.14309810e-01  -9.48885201e-01]
#  [ -1.54905600e+00  -3.75876244e-01   1.62187521e+00]
#  [  7.63388140e-01  -3.85305282e-01   3.48763075e-01]
#  [  2.39027306e+00   1.83219214e+00   6.23528107e-01]]
#
# [[ -3.53636246e+13   2.43113968e+13   1.41308765e+13]
#  [ -2.52635311e+13  -2.55198767e+12   7.25048576e+12]
#  [ -1.19041584e+13  -1.31181687e+13  -4.63829337e+12]
#  [  6.56647231e+12  -2.24892236e+13  -4.12367508e+11]]]

In this array you will access the momentum, I and (X,Y,Z) using the first, second and third indices:

a[0,1,2] #first momentum, second I and third coordinate