Gen_Electrondat.py

import MDAnalysis as mda
import argparse
def main():
    while True:
        whoami = """WRITE ME"""
        parser = argparse.ArgumentParser(description = whoami, formatter_class = argparse.RawDescriptionHelpFormatter)
        parser.add_argument('-tpr','--tpr',type=str, required=True, nargs=1, help='Name of tpr file')
        parser.add_argument('-lip', '--lip', required=True, action='append', help='Call -lip and provide the resname of each membrane lipid')
 

        args = parser.parse_args()

        tpr = args.tpr[0]
        lipids = args.lip
        break

    u = mda.Universe(tpr)
    lines = []
    for lip in lipids:
        res = u.select_atoms(f'resname {lip}').split('residue')
        names = res[0].names
        for n in names:
            if n[0] == 'H':
                lines.append(f'{n} = 1\n')
            
            if n[0] == 'C':
                lines.append(f'{n} = 6\n')

            if n[0] == 'N':
                lines.append(f'{n} = 7\n')

            if n[0] == 'O':
                lines.append(f'{n} = 8\n')

    with open('electron.dat', 'w') as file:
        file.write(f'{len(lines)}\n')
        for lin in lines:
            file.write(f'{lin}')


            


    return

if __name__ == '__main__':
    main()