extractsamples

#!/usr/bin/env python
import os
import numpy as np
import subprocess 
import molrecog 
import itertools
b2A = 0.529177249
A2b = 1./b2A

def head(iterator):
    return list(next(iterator))

def chunkedList(iterable, n):
    def inner_chunked():
        if n < 1:
            raise ValueError('n must be at least one')
        iterator = iter(iterable)

        while chunk := list(itertools.islice(iterator, n)):
            yield chunk
    return inner_chunked()

def getGeoms(cwd, nrAtoms):
    fileName = cwd + "/final_output"
    fileLines = lambda: (line for line in open(fileName,'r'))
    geomLines = lambda: (
        fileLine.strip().split() 
        for fileLine in fileLines()
        if len(fileLine.strip().split()) == 6
    ) 

    atomCoords = (
        [float(atomCoord) for atomCoord in geomLine[2:5]]
        for geomLine in geomLines()
    )

    atomNames = (
        geomLine[0] 
        for geomLine in geomLines()
    )

    geomCoords = chunkedList(atomCoords, nrAtoms)
    atomNames = head(chunkedList(atomNames, nrAtoms))
    atomNames = [
        atomName + str(iAtomName)
        for iAtomName, atomName in enumerate(atomNames)
    ]
    
    return atomNames, geomCoords

def main():
    cwd = os.getcwd()
    atomNames, geoms = getGeoms(cwd, 5)
    for iGeom, geom in enumerate(geoms):
        currDir = cwd + '/ICs/IC' + str(iGeom + 1)
        writer = molrecog.Writer(geom, atomNames, currDir)
        writer.write('bagel')

if __name__ == "__main__":
    main()