ConnectRes

#!/usr/bin/env python
# ConnectRes - Connect atoms in openMM pdb files and adds 
# force field parameters for these bonds, and other parameters
# into an openMM xml force field file.
# Ryan Clark <ryan.clark@ens-lyon.fr>
# Agilio Padua <agilio.padua@ens-lyon.fr>
# version 2021/05/14
# coding=utf8

import numpy as np
import argparse
import sys
from collections import Counter
import xml.etree.ElementTree as ET
from copy import deepcopy

#-------------
bondtol = 0.5
#-------------

def dist2at(at1,at2,boxDim):
	'''Find shortest distance between 2 atoms given periodic boundary conditions'''
	boxX = float(boxDim[0]);boxY = float(boxDim[1]);boxZ = float(boxDim[2])
	x = abs(at1['x']-at2['x'])
	while x > boxX/2:
		x = abs(x-boxX)
	y = abs(at1['y']-at2['y'])
	while y > boxY/2:
		y = abs(y-boxY)
	z = abs(at1['z']-at2['z'])
	while z > boxZ/2:
		z = abs(z-boxZ)
	return np.sqrt(x**2+y**2+z**2)

def dist2atNoWrap(at1,at2):
	'''Find shortest distance between 2 atoms given periodic boundary conditions'''
	x = abs(at1['x']-at2['x'])
	y = abs(at1['y']-at2['y'])
	z = abs(at1['z']-at2['z'])
	return np.sqrt(x**2+y**2+z**2)



#-----------------------

class PdbFile(object):
	'''Class object for pdb structure file for openMM'''
	def __init__(self, datafile, ff):
		residues = ff.root.find('Residues')
		atoms_ff = ff.root.find('AtomTypes')
		self.box = ''
		self.atoms = []
		self.bonds = []
		self.newConnections_created = [] # List for atom ids of new connections (needs creating when initialising)
		with open(datafile, 'r') as f:
			line = f.readline()
			while 'CRYST1' not in line:
				line = f.readline()
			self.box = line.strip()
			line = f.readline()
			while not (line.startswith('ATOM') or line.startswith('HETATM')):
				line = f.readline()
			while 'ATOM' in line or 'HETATM' in line:
				atom = {}
				atom['index'] = int(line[6:11])
				atom['name'] = line[12:16].strip()
				atom['frag'] = line[17:20].strip()
				atom['seq'] = int(line[22:26])
				atom['x'] = float(line[30:38])
				atom['y'] = float(line[38:46])
				atom['z'] = float(line[46:54])
				atom['ele'] = line[76:78].strip()
				# uses residues in force field file to assign atom class
				for residue in residues:
					if atom['frag'] in residue.get('name')[0:3]:
						for ffatom in residue.iter('Atom'):
							if ffatom.get('name') == atom['name']:
								atom['type'] = ffatom.get('type')
				for at in atoms_ff:
					if at.get('name') == atom['type']:
						atom['class'] = at.get('class')
				self.atoms.append(atom)
				line = f.readline()
			while 'CONECT' in line:
				i = int(line[6:11])
				j = int(line[11:16])
				if self.atoms[i-1]['seq'] != self.atoms[j-1]['seq']:
					print('WARNING: bond atoms {0:d} and {1:d} not in same molecule'.format(i, j))
					#sys.exit(1)
				self.bonds.append({'i': i, 'j': j})
				line = f.readline()	
		#Pull box dimensions out of header. Only supports orthogonal boxes so far
		box_line = self.box.split()
		if float(box_line[4]) != 90.0 or float(box_line[5]) != 90.0 or float(box_line[5]) != 90.0 :
			print('ERROR: Non-orthogonal boxes not supported')
			sys.exit(1)
		self.boxDimensions = [box_line[1],box_line[2],box_line[3]]

	def PrintToFile(self, outfile):
		'''Print .pdb file of PDB class'''
		print('Writing PDB file')
		with open(outfile, 'w') as f:
			f.write('TITLE     created by ConnectRes\n')
			f.write('REMARK    includes extra ff parameters and external bonds \n')
			f.write(self.box + '\n')
			n = 0
			for at in self.atoms:
				n += 1
				f.write('HETATM{0:5d} {1:4s} {2:3s}  {3:4d}    '\
				'{4:8.3f}{5:8.3f}{6:8.3f}  1.00  0.00          '\
				'{7:>2s}\n'.format(at['index'], at['name'], at['frag'], at['seq'],
				at['x'], at['y'], at['z'], at['ele']))
			for bd in self.bonds:
				f.write('CONECT{0:5d}{1:5d}\n'.format(bd['i'], bd['j']))
			f.write('END\n')

	def AddConnection(self, atom1, atom2, ff, toPDB, BondFile, AngleFile, DihedralFile):
		'''Add connection between 2 atoms'''
		atom1_class = atom1.split(',')[0]; atom1_frag = atom1.split(',')[1]
		atom2_class = atom2.split(',')[0]; atom2_frag = atom2.split(',')[1]
		atom1_names = []; atom2_names = [] # List of atom names that will appear in pdb file
		
		# Use force field bonds to find bond length for distance search
		bondtypes = ff.root.find('HarmonicBondForce')
		for bond in bondtypes:
			bondat1 = bond.get('class1'); bondat2 = bond.get('class2')
			if atom1_class in [bondat1,bondat2] and atom2_class in [bondat1,bondat2]:
				bond_length = float(bond.get('length'))*10.
		if 'bond_length' not in locals():
			print('ERROR: No bond between {0} and {1} specified in force field'\
			.format(atom_1_class,atom_2_class))
			sys.exit(1)  
		
		# Populate list of atom names for pdb search
		residues = ff.root.find('Residues')
		for residue in residues:
			res_name = residue.get('name')
			for ffatom in residue.iter('Atom'):
				ffatom_type = ffatom.get('type').split('-')[1]
				ffatom_name = ffatom.get('name')
				if ffatom_type == atom1_class and res_name == atom1_frag:
					atom1_names.append('{0}'.format(ffatom_name))
				if ffatom_type == atom2_class and res_name == atom2_frag:
					atom2_names.append('{0}'.format(ffatom_name))
		
		# Make lists of atoms of the correct type
		at_1_list = []; at_2_list = []
		for atom in self.atoms:
			if atom['name'] in atom1_names and atom['frag'] == atom1_frag and atom['name'][0] != 'D':
				at_1_list.append(atom)
			if atom['name'] in atom2_names and atom['frag'] == atom2_frag and atom['name'][0] != 'D':
				at_2_list.append(atom)
		
		# Run through lists of atoms and connect those close enough according to bond length
		for atom_1 in at_1_list:
			no_bonds = 0
			for atom_2 in at_2_list:
				distance = dist2at(atom_1, atom_2, self.boxDimensions)
				if distance<bond_length+bondtol:
					no_bonds+=1
					if toPDB==True:
						self.bonds.append({'i': atom_1['index'], 'j': atom_2['index']})
					else:
						BondFile.write('{0:5d}\t{1:5d}\n'.format(atom_1['index'],atom_2['index']))
					self.newConnections_created.append([atom_1, atom_2])
					
		if toPDB == False:
			print(' Bonds between {0} and {1} added to file {2} '.format(atom_1['class'], atom_2['class'],BondFile.name))
			
			angletypes = ff.root.find('HarmonicAngleForce')
			# Check every angle
			for ang in angletypes:
				angAt1 = ang.get('class1'); angAt2 = ang.get('class2'); angAt3 = ang.get('class3')
				angK = ang.get('k'); angLength = ang.get('angle')
				# If the angle has atoms which are a part of the bond which has just been added
				if atom_1['class'] in (angAt1,angAt2,angAt3) and atom_2['class'] in (angAt1,angAt2,angAt3):
					# Make lists of all of the atoms which match atom1, atom2, atom3 in the angle
					At1 = []; At2 = []; At3 = []
					
					for atom in self.atoms:
						if atom['class'] == angAt1 and atom['frag'] in (atom1_frag,atom2_frag) and atom['name'][0] != 'D':
							At1.append(atom)
						if atom['class'] == angAt2 and atom['frag'] in (atom1_frag,atom2_frag) and atom['name'][0] != 'D':
							At2.append(atom)
						if atom['class'] == angAt3 and atom['frag'] in (atom1_frag,atom2_frag) and atom['name'][0] != 'D':
							At3.append(atom)
					
					# if the distance between at1-at2-at3 are small, angle needs to be added, so add it
					for at1 in At1:
						for at2 in At2:
							if dist2at(at1,at2,self.boxDimensions) < 3:
								for at3 in At3:
									if dist2at(at2,at3,self.boxDimensions) < 3:
										AngleFile.write('{0:5d}\t{1:5d}\t{2:5d}\t{3}\t{4}\n'.format(at1['index'],at2['index'],at3['index'],angK,angLength))
				
			print(' Angles including atoms {0} and {1} added to file {2} '.format(atom_1['class'],atom_2['class'],AngleFile.name))
			
			##TODO: Currently takes all dihedrals on an element type (chemical element) and creates all dihedrals (i.e. Zn, NZ, CX, CZ makes all dihedrals with CX/CZ in column 3 and 4). This means it's putting CX and CZ into the last 2 atom lists in this case. Need to make specific to each atom type/class/name whichever is more appropriate. Can also add something to make sure that atoms are not repeated.
			
			# Make list of dihedrals
			dihedrals = []
			dihedraltypes = ff.root.find('RBTorsionForce')
			# Check every dihedral
			for dih in dihedraltypes:
				dihAt1 = dih.get('class1'); dihAt2 = dih.get('class2'); dihAt3 = dih.get('class3'); dihAt4 = dih.get('class4')
				c0 = dih.get('c0');c1 = dih.get('c1');c2 = dih.get('c2');c3 = dih.get('c3');c4 = dih.get('c4');c5 = dih.get('c5')
				# If the dihedral has atoms which are a part of the bond which has just been added
				if atom_1['class'] in (dihAt1,dihAt2,dihAt3,dihAt4) and atom_2['class'] in (dihAt1,dihAt2,dihAt3,dihAt4):
					# Make lists of all of the atoms which match atom1, atom2, atom3, atom4 in the dihedral
					At1 = []; At2 = []; At3 = []; At4 = []
					
					for atom in self.atoms:
						if atom['class'] == dihAt1 and atom['frag'] in (atom1_frag,atom2_frag) and atom['name'][0] != 'D':
							At1.append(atom)
						if atom['class'] == dihAt2 and atom['frag'] in (atom1_frag,atom2_frag) and atom['name'][0] != 'D':
							At2.append(atom)
						if atom['class'] == dihAt3 and atom['frag'] in (atom1_frag,atom2_frag) and atom['name'][0] != 'D':
							At3.append(atom)
						if atom['class'] == dihAt4 and atom['frag'] in (atom1_frag,atom2_frag) and atom['name'][0] != 'D':
							At4.append(atom)
							
					# if the distance between at1->at2->at3->at4 are small, dihedral needs to be added, so add it
					for at1 in At1:
						for at2 in At2:
							for bond in bondtypes:
								bondat1 = bond.get('class1'); bondat2 = bond.get('class2')
								if at1['class'] in [bondat1,bondat2] and at2['class'] in [bondat1,bondat2]:
									b_l = float(bond.get('length'))*10.
							if 0.1 < dist2at(at1,at2,self.boxDimensions) < b_l+bondtol:
								for at3 in At3:
									for bond in bondtypes:
										bondat1 = bond.get('class1'); bondat2 = bond.get('class2')
										if at2['class'] in [bondat1,bondat2] and at3['class'] in [bondat1,bondat2]:
											b_l = float(bond.get('length'))*10.
									if 0.1 < dist2at(at2,at3,self.boxDimensions) < b_l+bondtol:	
										for at4 in At4:
											for bond in bondtypes:
												bondat1 = bond.get('class1'); bondat2 = bond.get('class2')
												if at3['class'] in [bondat1,bondat2] and at4['class'] in [bondat1,bondat2]:
													b_l = float(bond.get('length'))*10.
											if 0.1< dist2at(at3,at4,self.boxDimensions) < b_l+bondtol:
												ats = [at1['index'],at2['index'],at3['index'],at4['index']]
												if at1['index'] != at3['index'] and at1['index'] != at4['index'] and at2['index'] != at4['index']:
													#print(at1,at2,at3,at4)
													ats += [c0,c1,c2,c3,c4,c5]	
													# if dihedral is not in list, add to list
													if ats not in dihedrals:
														dihedrals.append(ats)
			dihs = []	
			for d in dihedrals:
				if d not in dihs:
					dihs.append(d)
			for d in dihs:
				DihedralFile.write('{0:5d}\t{1:5d}\t{2:5d}\t{3:5d}\t{4}\t{5}\t{6}\t{7}\t{8}\t{9}\n'.format(d[0],d[1],d[2],d[3],d[4],d[5],d[6],d[7],d[8],d[9]))
				
			print(' Dihedrals including atoms {0} and {1} added to file {2} '.format(atom_1['class'],atom_2['class'],DihedralFile.name))
	
	def reportConnections(self):
		'''Print output report of connections created for each atom pair'''
		at_con = []; no_con = []; no_at_con = []
		
		# Check every connection
		for i in range(len(self.newConnections_created)):
			con = self.newConnections_created[i]
			new_con = [con[0]['name'],con[1]['name'],con[0]['frag'],con[1]['frag']]
			
			# Count connections between atom:residue pairs
			if new_con not in at_con:
				at_con.append(new_con)
				no_con.append(1)
			else:
				j=0;
				for j in range(len(at_con)):
					if at_con[j] == new_con:
						no_con[j]+=1
		
		# Print number of bonds creates per atom:residue pair
		for i in range(len(no_con)):
			print(' {0} Bonds created between atom:{1} in residue:{2} and atom:{3} in residue:{4}'.format(no_con[i], at_con[i][0], at_con[i][2], at_con[i][1], at_con[i][3]))

	def RemoveLongBonds(self):
		for bond in self.bonds:
			at1 = bond['i']
			at2 = bond['j']
			for atom in self.atoms:
				if atom['index'] == at1:
					atom1 = atom
				if atom['index'] == at2:
					atom2 = atom
			dist = dist2atNoWrap(atom1,atom2)
			#print(atom1,atom2)
			if dist>10:
				self.bonds.remove(bond)
				#print(atom1,atom2)
				#print(dist)


# ------------------------

def indent_xml(elem, level=0, hor=' ', ver='\n'):
	'''pretty-print xml tree'''
	
	spc = ver + level * hor
	if len(elem):
		if not elem.text or not elem.text.strip():
			elem.text = spc + hor
		if not elem.tail or not elem.tail.strip():
			elem.tail = spc
		for elem in elem:
			indent_xml(elem, level + 1, hor, ver)
		if not elem.tail or not elem.tail.strip():
			elem.tail = spc
	else:
		if level and (not elem.tail or not elem.tail.strip()):
			elem.tail = spc



# ---------------------

class XmlFile(object):
	'''Class object for xml force field files for openMM'''
	def __init__(self, infile):
		'''element tree from xml file'''
		self.ftree = ET.parse(infile)
		self.root = self.ftree.getroot()

	def PrintToFile(self, outfile):
		'''write force field to xml file'''
		print('Writing XML file')
		indent_xml(self.root)
		self.ftree.write(outfile)

	def AddParams(self, ff):
		'''Add params file data to xml'''
		for atom in ff.atoms:
			print('------------------------------------------------------------------\n'\
			'WARNING: ATOMS FOUND IN ADDITION FORCE FIELD FILE\nThe ability to add atoms '\
			'is not included here, all atoms\nshould be included in original fftool input.'\
			'\n------------------------------------------------------------------')
		
		for bond in ff.bonds:
			# Chack bond is of valid type
			if bond[2] != 'harm':
				print('ERROR: Currently only harmonic bonds are supported')
				sys.exit(1)
			
			# Load bond parameters from input .ff file
			at1 = bond[0]; at2 = bond[1]
			at1inFF = False; at2inFF = False
			length = bond[3][0]/10.; k = bond[3][1]*100.
			
			# Check atoms in bond are valid atoms in xml file
			atomtypes = self.root.find('AtomTypes')
			for atom in atomtypes:
				atom_class = atom.get('class')
				if at1 == atom_class:
					at1inFF = True
				if at2 == atom_class:
					at2inFF = True
			
			# If atoms are valid put into xml file. Print error if not.
			if at1inFF and at2inFF:
				bondtypes = self.root.find('HarmonicBondForce')
				newbond = ET.SubElement(bondtypes, 'Bond')
				newbond.set('class1', at1)
				newbond.set('class2', at2)
				newbond.set('length', '{0:.5f}'.format(length))
				newbond.set('k', '{0:.1f}'.format(k))
				print(' Bond added to force field: {0} - {1}'.format(at1,at2))
			elif not at1inFF:	
				print('ERROR: cannot find atom {0} in force field'.format(at1))
			elif not at2inFF:	
				print('ERROR: cannot find atom {0} in force field'.format(at2))
		
		for angle in ff.angles:
			# Check angle is of valid type
			if angle[3] != 'harm':
				print('ERROR: Currently only harmonic angles are supported')
				sys.exit(1)
			
			# Load angle parameters from input .ff file
			at1 = angle[0]; at2 = angle[1]; at3 = angle[2]
			at1inFF = False; at2inFF = False; at3inFF = False
			ang = angle[4][0] * np.pi/180.0; k = angle[4][1]
			
			# Check atoms in angle are valid atoms in xml file
			atomtypes = self.root.find('AtomTypes')
			for atom in atomtypes:
				atom_class = atom.get('class')
				if at1 == atom_class:
					at1inFF = True
				if at2 == atom_class:
					at2inFF = True
				if at3 == atom_class:
					at3inFF = True
			
			# If atoms are valid put into xml file. Print error if not.
			if at1inFF and at2inFF and at3inFF:
				angtypes = self.root.find('HarmonicAngleForce')
				newang = ET.SubElement(angtypes, 'Angle')
				newang.set('class1', at1)
				newang.set('class2', at2)
				newang.set('class3', at3)
				newang.set('angle', '{0:.5f}'.format(ang))
				newang.set('k', '{0:.2f}'.format(k))
				print(' Angle added to force field: {0} - {1} - {2}'.format(at1,at2,at3))
			elif not at1inFF:	
				print('ERROR: cannot find atom {0} in force field'.format(at1))
			elif not at2inFF:	
				print('ERROR: cannot find atom {0} in force field'.format(at2))
			elif not at3inFF:	
				print('ERROR: cannot find atom {0} in force field'.format(at3))
		
		for dihed in ff.dihedrals:
			# Check dihedral is of valid type
			if dihed[4] != 'opls':
				print('ERROR: Currently only opls dihedrals are supported')
				sys.exit(1)
			
			# Load dihedral parameters from input .ff file
			at1 = dihed[0];at2 = dihed[1]; at3 = dihed[2];at4 = dihed[3]
			at1inFF = False; at2inFF = False; at3inFF = False; at4inFF = False
			par0 = dihed[5][0]; par1 = dihed[5][1]; par2 = dihed[5][2]; par3 = dihed[5][3]
			eps = 1e-12
			#convert OPLS to RB -- from fftool
			c0 = par1 + 0.5 * (par0 + par2) + eps
			c1 = 0.5 * (-par0 + 3 * par2) + eps
			c2 = -par1 + 4 * par3 + eps
			c3 = -2 * par2 + eps; c4 = -4 * par3 + eps; c5 = 0.0
			
			# Check atoms in dihedral are valid atoms in xml file
			atomtypes = self.root.find('AtomTypes')
			for atom in atomtypes:
				atom_class = atom.get('class')
				if at1 == atom_class:
					at1inFF = True
				if at2 == atom_class:
					at2inFF = True
				if at3 == atom_class:
					at3inFF = True
				if at4 == atom_class:
					at4inFF = True
			
			# If atoms are valid put into xml file. Print error if not.
			if at1inFF and at2inFF and at3inFF and at4inFF:
				dihedtypes = self.root.find('RBTorsionForce')
				newdihed = ET.SubElement(dihedtypes, 'Proper')
				newdihed.set('c0', '{0:.5f}'.format(c0)); newdihed.set('c1', '{0:.5f}'.format(c1))
				newdihed.set('c2', '{0:.5f}'.format(c2)); newdihed.set('c3', '{0:.5f}'.format(c3))
				newdihed.set('c4', '{0:.5f}'.format(c4)); newdihed.set('c5', '{0:.5f}'.format(c5))
				newdihed.set('class1', at1); newdihed.set('class2', at2)
				newdihed.set('class3', at3); newdihed.set('class4', at4)
				print(' Dihedral added to force field: {0} - {1} - {2} - {3}'.format(at1,at2,at3,at4))
			elif not at1inFF:	
				print('ERROR: cannot find atom {0} in force field'.format(at1))
			elif not at2inFF:	
				print('ERROR: cannot find atom {0} in force field'.format(at2))
			elif not at3inFF:	
				print('ERROR: cannot find atom {0} in force field'.format(at3))
			elif not at4inFF:	
				print('ERROR: cannot find atom {0} in force field'.format(at4))
		
		for improp in ff.impropers:
			# Check improper is of valid type
			if improp[4] != 'opls':
				print('ERROR: Currently only opls improper dihedrals are supported')
				sys.exit(1)
			
			# Load improper parameters from input .ff file
			at1 = improp[0];at2 = improp[1]; at3 = improp[2];at4 = improp[3]
			at1inFF = False; at2inFF = False; at3inFF = False; at4inFF = False
			par0 = improp[5][0]; par1 = improp[5][1]; par2 = improp[5][2]; par3 = improp[5][3]
			eps = 1e-12
			#convert OPLS to RB -- from fftool
			c0 = par1 + 0.5 * (par0 + par2) + eps
			c1 = 0.5 * (-par0 + 3 * par2) + eps
			c2 = -par1 + 4 * par3 + eps
			c3 = -2 * par2 + eps; c4 = -4 * par3 + eps; c5 = 0.0
			
			# Check atoms in improper are valid atoms in xml file
			atomtypes = self.root.find('AtomTypes')
			for atom in atomtypes:
				atom_class = atom.get('class')
				if at1 == atom_class:
					at1inFF = True
				if at2 == atom_class:
					at2inFF = True
				if at3 == atom_class:
					at3inFF = True
				if at4 == atom_class:
					at4inFF = True
			
			# If atoms are valid put into xml file. Print error if not.
			if at1inFF and at2inFF and at3inFF and at4inFF:
				dihedtypes = self.root.find('RBTorsionForce')
				newdihed = ET.SubElement(dihedtypes, 'Improper')
				newdihed.set('c0', '{0:.5f}'.format(c0)); newdihed.set('c1', '{0:.5f}'.format(c1))
				newdihed.set('c2', '{0:.5f}'.format(c2)); newdihed.set('c3', '{0:.5f}'.format(c3))
				newdihed.set('c4', '{0:.5f}'.format(c4)); newdihed.set('c5', '{0:.5f}'.format(c5))
				newdihed.set('class1', at1); newdihed.set('class2', at2)
				newdihed.set('class3', at3); newdihed.set('class4', at4)
				print(' Improper dihedral added to force field: {0} - {1} - {2} - {3}'.format(at1,at2,at3,at4))
			elif not at1inFF:	
				print('ERROR: cannot find atom {0} in force field'.format(at1))
			elif not at2inFF:	
				print('ERROR: cannot find atom {0} in force field'.format(at2))
			elif not at3inFF:	
				print('ERROR: cannot find atom {0} in force field'.format(at3))
			elif not at4inFF:	
				print('ERROR: cannot find atom {0} in force field'.format(at4))
			
		

	def AddExternalBonds(self, pdb):
		'''Add external bonds to xml'''
		bondsCreated = pdb.newConnections_created
		# List of atoms and number of times they appear in new bond list 
		ext_bond_atom_list = []
		atom_list = []
		count_list = []
		
		for bond in bondsCreated:
			for atom in bond:
				at_list = [atom['index'],atom['name'],atom['frag']]
				if at_list not in atom_list:
					atom_list.append(at_list)
					count_list.append(1)
				else:
					for i in range(len(atom_list)):
						if atom_list[i] == at_list:
							count_list[i]+=1
		
		# Use this list to find number of external bonds per atom:residue pair
		ext = []; ext_check = []
		for i in range(len(atom_list)):
			at = atom_list[i][1]
			frag = atom_list[i][2]
			count = count_list[i]
			at_ext = [at, frag, count]
			if at_ext not in ext:
				ext.append(at_ext)
				ext_check.append([at, frag])
		
		# Check number of times each atom:res pair appears in list, if more than one then print error and exit	
		for i in range(len(ext_check)):
			ext_count = ext_check.count(ext_check[i])
			if ext_count>1:
				print('ERROR: Atom {0} in residue {1} has multiple ExternalBonds of different types'.format(ext_check[i][0],ext_check[i][1]))
				sys.exit(1)
		
		# Check residues for atoms and residues that match those that need external bonds and add external bonds to residues 
		print('Adding external bond flags')
		residues = self.root.find('Residues')
		for ext_it in ext:
			at = ext_it[0]
			frag = ext_it[1]
			nobonds = ext_it[2]
			for residue in residues:
				res_name = residue.get('name')
				if res_name == frag:
					for ffatom in residue.iter('Atom'):
						if ffatom.get('name') == at:
							for i in range(nobonds):	
								newExt = ET.SubElement(residue, 'ExternalBond')
								newExt.set('atomName', at)
							print(' Added {0} ExternalBond flags added for atom {1} in residue {2}'.format(nobonds, at, frag))
						
					
				
			
		



#--------------------

class Params(object):	
	''' Parameters in force field file. init from fftool.'''
	def __init__(self, filename):
		
		self.filename = filename
		self.atoms = []
		self.bonds = []
		self.angles = []
		self.dihedrals = []
		self.impropers = []
		for line in open(filename, 'r'):
			if line.startswith('#') or line.strip() == '':
				continue
			if line.lower().startswith('atom'):
				section = 'atoms'
				continue
			elif line.lower().startswith('bond'):
				section = 'bonds'
				continue
			elif line.lower().startswith('angl'):
				section = 'angles'
				continue
			elif line.lower().startswith('dihe'):
				section = 'dihedrals'
				continue
			elif line.lower().startswith('impro'):
				section = 'improper'
				continue
			tok = line.strip().split()
			if section == 'atoms':
				name = tok[0]
				attp = tok[1]
				m = float(tok[2])
				q = float(tok[3])
				pot = tok[4]
				par = [float(p) for p in tok[5:]]
				self.atoms.append([name, attp, m, q, pot, par])
			elif section == 'bonds':
				iatp = tok[0]
				jatp = tok[1]
				pot = tok[2]
				par = [float(p) for p in tok[3:]]
				self.bonds.append([iatp, jatp, pot, par])
			elif section == 'angles':
				iatp = tok[0]
				jatp = tok[1]
				katp = tok[2]
				pot = tok[3]
				par = [float(p) for p in tok[4:]]
				self.angles.append([iatp, jatp, katp, pot, par])
			elif section == 'dihedrals':
				iatp = tok[0]
				jatp = tok[1]
				katp = tok[2]
				latp = tok[3]
				pot = tok[4]
				par = [float(p) for p in tok[5:]]
				self.dihedrals.append([iatp, jatp, katp, latp, pot, par])
			elif section == 'improper':
				iatp = tok[0]
				jatp = tok[1]
				katp = tok[2]
				latp = tok[3]
				pot = tok[4]
				par = [float(p) for p in tok[5:]]
				self.impropers.append([iatp, jatp, katp, latp, pot, par])



#----------------

def main():
	parser = argparse.ArgumentParser(
		description = 'Take a .pdb file with residues created using fftool and connects the specified atoms across residues in the .pdb structure file (with connect commands). Will also add bond, angle, and dihedral data to the .xml force field file from the input force field. This includes adding external bond flags between residues.\n NOTE: Does not import atoms or LJ parameters in force field file, just reads these in and outputs as they appear in original force field. Therefore these cannot be edited in force fields.')
	parser.add_argument('fragments', nargs='+',
						help = 'atom1,frag1 atom2,frag2 [atom3,frag3 atom4,frag4].... Pairs of atoms to bond together across fragments')
	parser.add_argument('-f', '--fffile', default = 'frag.ff',
						help = 'Force field file containing bond, angle and dihedral data to add to .xml file ') 
	parser.add_argument('-p', '--conectinpdb', action = 'store_true', 
						help = 'Print connect to pdb file. NOTE: It is not recommended to do this if the bonds you create cross periodic boundary conditions!! If not specified, bonds, angles and dihedrals that need to be added to your openMM system are printed to files "Bonds.txt", "Angles.txt" and "Dihedrals.txt".') 
	parser.add_argument('-id', '--indfile', default = 'config.pdb',
						help = 'input openMM data file (default: config.pdb)')  
	parser.add_argument('-od', '--outdfile', default = 'config-c.pdb',
						help = 'output openMM data file (default: config-c.pdb)')
	parser.add_argument('-if', '--inffile', default = 'field.xml',
						help = 'input openMM force field file (default: field.xml)')  
	parser.add_argument('-of', '--outffile', default = 'field-c.xml',
						help = 'output openMM force field file (default: field-c.xml)')    
	args = parser.parse_args()
	
	if len(args.fragments)%2 != 0:
		print('ERROR: Must specify fragments in pairs')
		sys.exit(1)
	else:
		nopairs = int(len(args.fragments)/2)
		print(' {0} pair(s) of atoms.'.format(nopairs))
		atoms_1 = [None]*nopairs; atoms_2 = [None]*nopairs
		frags_1 = [None]*nopairs; frags_2 = [None]*nopairs
		for i in range(nopairs):
			atoms_1[i],frags_1[i] = args.fragments[2*i].split(',')
			atoms_2[i],frags_2[i] = args.fragments[2*i+1].split(',')
			print('  Linking {0} in fragment {1} to {2} in fragment {3}'\
			.format(atoms_1[i],frags_1[i],atoms_2[i],frags_2[i]))
	
	if args.conectinpdb:
		toPDB = True
		BondFile = None
		AngleFile = None
		DihedralFile = None
	else:
		toPDB = False
		BondFile = open('Bonds.txt','w')
		AngleFile = open('Angles.txt','w')
		DihedralFile = open('Dihedrals.txt','w')
	
	print('Reading force field and extra data files')
	ForceField = XmlFile(args.inffile)
	ParametersToAdd = Params(args.fffile)
	
	print('Reading pdb file')
	Atoms = PdbFile(args.indfile, ForceField)
	
	print('Adding new parameters to force field file')
	ForceField.AddParams(ParametersToAdd)
	
	print('Adding connections and external bonds')
	for i in range(nopairs):
		Atoms.AddConnection(','.join([atoms_1[i],frags_1[i]]), ','.join([atoms_2[i],frags_2[i]]), ForceField, toPDB, BondFile, AngleFile, DihedralFile)
	Atoms.reportConnections()
	ForceField.AddExternalBonds(Atoms)
	
	print('\nNOTE: The ignoreExternalBond flag may be need to included during system creation.')
	print(' system = forcefield.createSystem(..., ignoreExternalBonds=True)')
	print(' Sometimes openMM does not like bonds over periodic boundary conditions.')
	print(' If so, then a CustomBondForce with setUsesPeriodicBoundaryConditions(True) may be needed.\n')
	
	#Atoms.RemoveLongBonds()
	
	if not args.conectinpdb:
		BondFile.close()
		AngleFile.close()
		DihedralFile.close()
		print('Connections across external bonds not added to PDB file.')
		print(' Bonds, angles and dihedrals will need to be added manually from files:\n {0}\n {1}\n {2}\n'.format(BondFile.name,AngleFile.name,DihedralFile.name))
		print('The following is an example loop to add bonds/angles from file to the openMM system after creation\n\nBonds = openmm.HarmonicBondForce()\n\nBonds.setUsesPeriodicBoundaryConditions(True)\nbondsToAdd = "Bonds.txt"\nk = 10000.0*unit.kilojoule/(unit.nanometer**2*unit.mole)\nr0 = 0.15*unit.nanometer\nfor line in open(bondsToAdd,"r"):\n\ttok = line.split()\n\tBonds.addBond(int(tok[0]),int(tok[1]),float(r0),float(k))\nsystem.addForce(Bonds)')
	
	Atoms.PrintToFile(args.outdfile)
	ForceField.PrintToFile(args.outffile)



# --------------------------

if __name__ == '__main__':
    main()