coulttxml

#!/usr/bin/env python
# Create code to add TT damping between Drude dipoles in OpenMM
# Vinicius Piccoli, Agilio Padua <agilio.padua@ens-lyon.fr>, 2024/07/17

import sys
import argparse
import xml.etree.ElementTree as ET

def intfromhexiflarge(a):
    """
    Convert string containing hex number (large systems) to int.
    VMD, OpenMM convention in PDB files with natom > 99999, nmol > 9999
    """

    if a.strip().isdigit():
        return int(a)
    
    ndig = len(a)
    declim = 10**ndig
    hexlim = int('A' + (ndig-1)*'0', 16)

    i = int(a, 16)
    return i - hexlim + declim


def count_molecules_from_pdb(pdb_filename):
    '''Count molecules in PDB file'''

    molecule_counts = {}
    with open(pdb_filename, 'r') as file:
        for line in file:
            if line.startswith('HETATM'):
                residue_name = line[17:20].strip()
                molecule_number = intfromhexiflarge(line[22:26])
                if residue_name not in molecule_counts:
                    molecule_counts[residue_name] = molecule_number
                else:
                    if molecule_number > molecule_counts[residue_name]:
                        molecule_counts[residue_name] = molecule_number
    numbers = [molecule_counts[res] for res in molecule_counts]
    for i in range(len(numbers) - 1, 0, -1):
        numbers[i] = numbers[i] - numbers[i-1]
    return numbers

def count_molecules_from_pdbx(pdbx_filename):
    '''Count molecules in PDBx/mmcif file'''

    molecule_counts = {}

    # First read to understand the structure of the PDBx/mmcif file
    with open(pdbx_filename, 'r') as file:
        columns = []
        resname = None
        nres = None
        for line in file:
        # Read the available categories in the PDBx/mmcif and store them in a "table of contents"
            if line.startswith('_atom'):
                fields = line.strip().split()
                columns.append(fields[0].strip())
        # Find indexes of columns with residue name and residue number
        for j in range (0,len(columns)):
            if columns[j] == '_atom_site.label_comp_id':
                resname = int(j)
            if columns[j] == '_atom_site.label_seq_id':
                nres = int(j)
        if resname == None:
            raise RuntimeError("Expected '_atom_site.label_comp_id' data in PDBx/mmcif config file")
        if nres == None:
            raise RuntimeError("Expected '_atom_site.label_seq_id' data in PDBx/mmcif config file")
    
    # Second read to pick data of interest 
    with open(pdbx_filename, 'r') as file:
        for line in file:
            fields = line.strip().split()
            if len(fields) == len(columns):
                residue_name = fields[resname].strip()
                molecule_number = int(fields[nres])
                if residue_name not in molecule_counts:
                    molecule_counts[residue_name] = 1
                else:
                    if molecule_number > molecule_counts[residue_name]:
                        molecule_counts[residue_name] = molecule_number
    numbers = [molecule_counts[res] for res in molecule_counts]
    for i in range(len(numbers) - 1, 0, -1):
        numbers[i] = numbers[i] - numbers[i-1]
    return numbers

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

before = '''
# OpenMM code to add charge-dipole damping in polarizable simulations

# CoulTT damping function
condition = "(don1*acc2 + don2*acc1) *"
coulttexp = condition + "q1*q2/(fpe0*r)*(- c*exp(-b*r) * (1 + b*r + (b*r)^2/2 + (b*r)^3/6 + (b*r)^4/24))"
CoulTT = openmm.CustomNonbondedForce(coulttexp)
CoulTT.setNonbondedMethod(openmm.NonbondedForce.CutoffPeriodic)
CoulTT.setCutoffDistance(1.2*unit.nanometer)
CoulTT.addGlobalParameter("fpe0", 0.007197587)
CoulTT.addGlobalParameter("c", 1.0)
CoulTT.addGlobalParameter("b", 52.0)
CoulTT.addPerParticleParameter("don")
CoulTT.addPerParticleParameter("acc")
CoulTT.addPerParticleParameter("q")

'''

after = '''
for i, f in enumerate(system.getForces()):
    f.setForceGroup(i)
    print(i, f.getName())
    if f.getName() == "NonbondedForce":
        nbf = f
for n in range(nbf.getNumExceptions()):
    ex = nbf.getExceptionParameters(n)
    CoulTT.addExclusion(ex[0], ex[1])
system.addForce(CoulTT)

'''

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

def main():
    parser = argparse.ArgumentParser(
        description = 'Write OpenMM code for CoulTT damping of charge-induced dipole interactions')
    parser.add_argument('--xml', default = 'field-p.xml', metavar = 'field-p.xml',
                        help = 'xml file with force field')
    parser.add_argument('--pdb', default = 'config-p.pdb', metavar = 'config-p.pdb',
                        help = 'PDB file with configuration')
    parser.add_argument('--core', action = 'store_true',
                        help = 'Use core charge instead of induced dipole charge in TT damping')
    args = parser.parse_args()

    print("force field from", args.xml)
    tree = ET.parse(args.xml)
    root = tree.getroot()

    residues_xml = root.find('Residues').findall('Residue')
    residue_names_xml = [residue.attrib['name'] for residue in residues_xml]


    if 'pdb' in args.pdb:
        num_molecules = count_molecules_from_pdb(args.pdb)
    elif 'mmcif' in args.pdb:
        num_molecules = count_molecules_from_pdbx(args.pdb)

    # Check if the number of residues matches the number of types in the PDB or PDBx file
    if len(residue_names_xml) != len(num_molecules):
        print("Error: The number of different residues in the XML does not match the count from the PDB or PDBx file.")
        sys.exit(1)

    with open('addCoulTT.py', 'w') as f:
        f.write(before)
        # Loop over each residue and its corresponding count from the PDB
        f.write("# Put flag 1 in 1st or 2nd field in atoms to be screened\n")
        f.write("# donor (1st field) is the charge, acceptor (2nd field) the induced dipole\n")
        for residue_xml, count in zip(residue_names_xml, num_molecules):
            f.write(f"for i in range({count}):\n")
            atoms = next((res for res in residues_xml if res.attrib['name'] == residue_xml), None).findall('Atom')
            for i, atom in enumerate(atoms):
                name = atom.attrib['name']
                charge = float(atom.attrib['charge'])
                if not args.core and i < len(atoms) - 1:
                    nxtname = atoms[i+1].attrib['name']
                    if nxtname.startswith('D'):
                        charge = - float(atoms[i+1].attrib['charge'])
                f.write(f"    CoulTT.addParticle([0, 0, {charge:8.5f}]) # {name:4s} {atom.attrib['type']:12s}\n")
        f.write(after)

    print("OpenMM code for CoulTT written to addCoulTT.py")

if __name__ == '__main__':
    main()