polxml

#!/usr/bin/env python
# polxml.py - add Drude oscillators to OpenMM xml force field file.
# Agilio Padua <agilio.padua@ens-lyon.fr>
# version 2024/02/08

import argparse
import random
import xml.etree.ElementTree as ET
import numpy as np


usage = """
==============================================================================
Add Drude oscillators to OpenMM xml force field file and pdb or pdbx/mmcif.
------------------------------------------------------------------------------
Format of file containing specification of Drude oscillators (alpha.ff):
  # type  dm/u  dq/e  k/(kJ/molA2)  alpha/A3  thole
  C3H     1.0   0.0   4184.0        1.016     2.6
  ...
* dm is the mass to place on the Drude particle (subtracted from its core),
* dq is the charge to place on the Drude particle (subtracted from its core),
* k is the harmonic force constant of the bond between core and Drude,
* alpha is the polarizability, hyrdogen aroms are not merged,
* thole is a parameter of the Thole damping function.
------------------------------------------------------------------------------
A Drude particle is created for each atom in the xml file
that corresponds to an atom type given in the Drude file.
This script will add new atom types, new bond types, new atoms and
new bonds to the xml file and new particles to the pdb file.
==============================================================================
"""

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

def hexiflarge(i, ndig):
    """
    Convert number to hex starting with A if more than n digits.
    VMD, OpenMM convention in PDB files with natom > 99999, nmol > 9999
    """

    declim = 10**ndig
    hexlim = int('A' + (ndig-1)*'0', 16)

    if i < declim:
        return f'{i}'
    else:
        i = i - declim + hexlim
        return f'{i:X}'

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

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

atomic_wt = {'H': 1.008, 'Li': 6.941, 'B': 10.811, 'C': 12.011,
             'N': 14.006, 'O': 15.999, 'F': 18.998, 'Ne': 20.180,
             'Na': 22.990, 'Mg': 24.305, 'Al': 26.982, 'Si':  28.086,
             'P': 30.974, 'S': 32.065, 'Cl': 35.453, 'Ar': 39.948,
             'K': 39.098, 'Ca': 40.078, 'Ti': 47.867, 'Fe': 55.845,
             'Cu': 63.546, 'Zn': 65.38, 'Se': 78.971, 'Br': 79.904, 'Kr': 83.798,
             'Mo': 95.96, 'Ru': 101.07, 'Sn': 118.710, 'Te': 127.60,
             'I': 126.904, 'Xe': 131.293}

def atomic_symbol(name):
    if name[:2] in atomic_wt:
        return name[:2]
    elif name[0] in atomic_wt:
        return name[0]
    else:
        print('warning: unknown symbol for atom ' + name)
        return name

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

def readpdb(Topology, pdbfile):
    '''read topology from PDB'''
    with open(pdbfile, 'r') as f:
        line = f.readline()
        while 'CRYST1' not in line:
            line = f.readline()
        Topology.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['n'] = intfromhexiflarge(line[6:11])
            atom['name'] = line[12:16].strip()
            atom['mol'] = line[17:20].strip()
            atom['seq'] = intfromhexiflarge(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()
            atom['nH'] = 0         # number of bonded H
            atom['type'] = ''      # non-bonded type
            atom['class'] = ''     # bonded type
            Topology.atoms.append(atom)
            line = f.readline()
        while 'CONECT' in line:
            i = intfromhexiflarge(line[6:11]) - 1
            j = intfromhexiflarge(line[11:16]) - 1
            if Topology.atoms[i]['seq'] != Topology.atoms[j]['seq']:
                print(f'bond between atoms {i:d} and {j:d} of different molecules')
                exit(1)
            Topology.bonds.append({'i': i, 'j': j})
            line = f.readline()
        Topology.natom = len(Topology.atoms)
    
def readpdbx(Topology, pdbxfile):
    '''read topology from PDBx/mmcif'''
    with open(pdbxfile, 'r') as f:
        line = f.readline()
        atoms_columns = []
        bonds_columns = []
            
        # Storage of columns numbers
        site_id = None
        type_symbol = None
        atom_id = None
        comp_id = None
        seq_id = None
        asym_id = None
        Cartn_x = None
        Cartn_y = None
        Cartn_z = None
        formal_charge = None

        ptnr1_label_seq_id = None
        ptnr2_label_seq_id = None
        ptnr1_label_atom_id = None
        ptnr2_label_atom_id = None

        for line in f:
        # Read available categories in PDBx/mmcif file and store them in a "table of contents"
            if line.startswith('_atom'):
                fields = line.strip().split()
                atoms_columns.append(fields[0].strip())
            if line.startswith('_struct'):
                fields = line.strip().split()
                bonds_columns.append(fields[0].strip())

        # Atoms section
        for j in range (0,len(atoms_columns)):
            if atoms_columns[j] == '_atom_site.id':
                site_id = int(j)
            if atoms_columns[j] == '_atom_site.type_symbol':
                type_symbol = int(j)
            if atoms_columns[j] == '_atom_site.label_atom_id':
                atom_id = int(j)
            if atoms_columns[j] == '_atom_site.label_comp_id':
                comp_id = int(j)
            if atoms_columns[j] == '_atom_site.label_seq_id':
                seq_id = int(j)
            if atoms_columns[j] == '_atom_site.label_asym_id':
                asym_id = int(j)
            if atoms_columns[j] == '_atom_site.Cartn_x':
                Cartn_x = int(j)
            if atoms_columns[j] == '_atom_site.Cartn_y':
                Cartn_y = int(j)
            if atoms_columns[j] == '_atom_site.Cartn_z':
                Cartn_z = int(j)
            if atoms_columns[j] == '_atom_site.pdbx_formal_charge':
                formal_charge = int(j)
        if site_id == None:
            raise RuntimeError("Expected '_atom_site.id' data in PDBx/mmcif config file")
        if type_symbol == None:
            raise RuntimeError("Expected '_atom_site.type_symbol' data in PDBx/mmcif config file")
        if atom_id == None:
            raise RuntimeError("Expected '_atom_site.label_atom_id' data in PDBx/mmcif config file")
        if comp_id == None:
            raise RuntimeError("Expected '_atom_site.label_comp_id' data in PDBx/mmcif config file")
        if seq_id == None:
            raise RuntimeError("Expected '_atom_site.label_seq_id' data in PDBx/mmcif config file")
        if asym_id == None:
            raise RuntimeError("Expected '_atom_site.label_asym_idl' data in PDBx/mmcif config file")
        if Cartn_x == None:
            raise RuntimeError("Expected '_atom_site.Cartn_x' data in PDBx/mmcif config file")
        if Cartn_y == None:
            raise RuntimeError("Expected '_atom_site.Cartn_y' data in PDBx/mmcif config file")
        if Cartn_z == None:
            raise RuntimeError("Expected '_atom_site.Cartn_z' data in PDBx/mmcif config file")
        if formal_charge == None:
            raise RuntimeError("Expected '_atom_site.pdbx_formal_charge' data in PDBx/mmcif config file")
            
        # Bonds section
        for j in range (0,len(bonds_columns)):
            if bonds_columns[j] == '_struct_conn.ptnr1_label_seq_id':
                ptnr1_label_seq_id = int(j)
            if bonds_columns[j] == '_struct_conn.ptnr2_label_seq_id':
                ptnr2_label_seq_id = int(j)
            if bonds_columns[j] == '_struct_conn.ptnr1_label_atom_id':
                ptnr1_label_atom_id = int(j)
            if bonds_columns[j] == '_struct_conn.ptnr2_label_atom_id':
                ptnr2_label_atom_id = int(j)
        if ptnr1_label_seq_id == None:
            raise RuntimeError("Expected '_struct_conn.ptnr1_label_seq_id' data in PDBx/mmcif config file")
        if ptnr2_label_seq_id == None:
            raise RuntimeError("Expected '_struct_conn.ptnr2_label_seq_id' data in PDBx/mmcif config file")
        if ptnr1_label_atom_id == None:
            raise RuntimeError("Expected '_struct_conn.ptnr1_label_atom_id' data in PDBx/mmcif config file")
        if ptnr2_label_atom_id == None:
            raise RuntimeError("Expected '_struct_conn.ptnr2_label_atom_id' data in PDBx/mmcif config file")
        
    # Second read to pick data of interest 
    with open(pdbxfile, 'r') as file:
        line = file.readline()
        for line in file:
            fields = line.strip().split()

            if line.startswith("_cell.length_a"):
                Topology.box.append(fields[1].strip())
            elif line.startswith("_cell.length_b"):
                Topology.box.append(fields[1].strip())
            elif line.startswith("_cell.length_c"):
                Topology.box.append(fields[1].strip())
            elif line.startswith("_cell.angle_alpha"):
                Topology.box.append(fields[1].strip())
            elif line.startswith("_cell.angle_beta"):
                Topology.box.append(fields[1].strip())
            elif line.startswith("_cell.angle_gamma"):
                Topology.box.append(fields[1].strip())
                
            if len(fields) == len(atoms_columns):
                atom = {}
                atom['n'] = int(fields[site_id])
                atom['name'] = fields[atom_id].strip()
                atom["mol"] = fields[comp_id].strip()
                atom['seq'] = int(fields[seq_id])
                atom['chain'] = int(fields[asym_id].strip())
                atom['x'] = float(fields[Cartn_x])
                atom['y'] = float(fields[Cartn_y])
                atom['z'] = float(fields[Cartn_z])
                atom['ele'] = fields[type_symbol].strip()
                atom['charge'] = float(fields[formal_charge])
                atom['nH'] = 0      # number of bonded H
                atom['type'] = ''   # non-bonded type
                atom['class'] = ''  # bonded type
                Topology.atoms.append(atom)
                
            if len(fields) == len(bonds_columns):
                atom_i = [int(fields[ptnr1_label_seq_id]), fields[ptnr1_label_atom_id].strip()]
                atom_j = [int(fields[ptnr2_label_seq_id]), fields[ptnr2_label_atom_id].strip()]
                i = None
                j = None
                for k in range (0, len(Topology.atoms)):
                    if Topology.atoms[k]['seq'] == atom_i[0]:
                        if Topology.atoms[k]['name'] == atom_i[1]:
                            i = k
                    if Topology.atoms[k]['seq'] == atom_j[0]:
                        if Topology.atoms[k]['name'] == atom_j[1]:
                            j = k
                Topology.bonds.append({'i': i, 'j': j})
    Topology.natom = len(Topology.atoms)

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

class Topology(object):
    '''topology and coordinates from pdb or pdbx/mmcif file'''

    def __init__(self, pdbfile):
        '''read from pdbx file'''

        self.box = []
        self.atoms = []
        self.bonds = []
        self.natom = 0
        self.ndrude = 0
        
        if pdbfile == 'config.pdb':
            readpdb(self, pdbfile)

        if pdbfile == 'config.mmcif':
            readpdbx(self, pdbfile)

    def writepdb(self, pdbfile):
        '''write topology to PDB file'''
        with open(pdbfile, 'w') as f:
            f.write('TITLE     created by polxml\n')
            f.write('REMARK    includes Drude particles\n')
            f.write(self.box + '\n')
            n = 0
            for at in self.atoms:
                n += 1
                atidx = hexiflarge(n, 5)
                molidx = hexiflarge(at['seq'], 4)
                f.write('HETATM{0:>5s} {1:4s} {2:3s}  {3:>4s}    '\
                        '{4:8.3f}{5:8.3f}{6:8.3f}  1.00  0.00          '\
                        '{7:>2s}\n'.format(atidx, at['name'][:4], at['mol'], molidx,
                        at['x'], at['y'], at['z'], at['ele']))
                if len(at['name']) > 4:
                    print('warning: atom name {} too long in pdb'.format(at['name']))
            for bd in self.bonds:
                atidx_i = hexiflarge(bd['i'] + 1, 5)
                atidx_j = hexiflarge(bd['j'] + 1, 5)
                f.write('CONECT{0:>5s}{1:>5s}\n'.format(atidx_i, atidx_j))
            f.write('END\n')

    def writepdbx(self, pdbxfile):
        #write topology to pdbx/mmcif file
        with open(pdbxfile, 'w') as f:
            f.write('data_\n')
            f.write('_cell.entry_id          simbox\n')
            f.write(f'_cell.length_a          {self.box[0]}\n')
            f.write(f'_cell.length_b          {self.box[1]}\n')
            f.write(f'_cell.length_c          {self.box[2]}\n')
            f.write(f'_cell.angle_alpha       {self.box[3]}\n')
            f.write(f'_cell.angle_beta        {self.box[4]}\n')
            f.write(f'_cell.angle_gamma       {self.box[5]}\n')
            f.write('loop_\n')
            f.write('_atom_site.group_PDB\n')
            f.write('_atom_site.id\n')
            f.write('_atom_site.type_symbol\n')
            f.write('_atom_site.label_atom_id\n')
            f.write('_atom_site.label_comp_id\n')
            f.write('_atom_site.label_seq_id\n')
            f.write('_atom_site.label_asym_id\n')
            f.write('_atom_site.Cartn_x\n')
            f.write('_atom_site.Cartn_y\n')
            f.write('_atom_site.Cartn_z\n')
            # Create a list to link atom names in a residue and the id (= index) of the residue. 
            atom_in_residue = []
            n = 0
            chain = 0
            for at in self.atoms:
                n += 1
                f.write('HETATM {0:8d} {1:2s} '\
                        '{2:5s} {3:4s} {4:7d} {5:2d} '\
                        '{6:12.5e} {7:12.5e} {8:12.5e}\n'.format(n, at['ele'], at['name'], at['mol'], at['seq'], chain, at['x'], at['y'], at['z']))
                atom_in_residue.append([at['name'], at['seq']])
            f.write('#\n')        
            f.write('loop_\n')
            f.write('_struct_conn.id\n')
            f.write('_struct_conn.ptnr1_label_seq_id\n')
            f.write('_struct_conn.ptnr2_label_seq_id\n')
            f.write('_struct_conn.ptnr1_label_atom_id\n')
            f.write('_struct_conn.ptnr2_label_atom_id\n')
            f.write('_struct_conn.ptnr1_label_asym_id\n')
            f.write('_struct_conn.ptnr2_label_asym_id\n')
            f.write('_struct_conn.conn_type_id\n')
            for bd in self.bonds:
                i = bd['i']
                j = bd['j']
                label_seq_id_i = self.atoms[i]['seq']
                label_seq_id_j = self.atoms[j]['seq']
                label_atom_id_i = self.atoms[i]['name']
                label_atom_id_j = self.atoms[j]['name']
                if i < j:
                    f.write(f'CONECT {label_seq_id_i:8d} {label_seq_id_j:8d} '
                            f'{label_atom_id_i:5s} {label_atom_id_j:5s} 0 0 covale\n')
                else:
                    f.write(f'CONECT {label_seq_id_j:8d} {label_seq_id_i:8d} '
                            f'{label_atom_id_j:5s} {label_atom_id_i:5s} 0 0 covale\n')
            f.write('data_\n')

    def settypes(self, ff):
        '''assign types from ff to atoms in topology'''

        for at in self.atoms:
            for res in ff.root.iter('Residue'):
                resname = res.get('name').replace('-', '').replace('+', '')[:3]
                if at['mol'] == resname:
                    for ffatom in res.iter('Atom'):
                        if at['name'] == ffatom.get('name'):
                            at['type'] = ffatom.get('type').split('-')[1]
                            for attype in ff.root.iter('Type'):
                                if at['type'] == attype.get('name').split('-')[1]:
                                    at['class'] = attype.get('class')
                                    break
                            break
                    break

    def countbondedH(self):
        '''count H bonded to each atom (to merge polarizabilities)'''

        for bd in self.bonds:
            i = bd['i']
            j = bd['j']
            if self.atoms[i]['name'][0] in 'Hh':
                self.atoms[j]['nH'] += 1
            elif self.atoms[j]['name'][0] in 'Hh':
                self.atoms[i]['nH'] += 1

    def checknH(self):
        '''check if each non-bonded type has same number of bonded H'''

        attype = [at['type'] for at in self.atoms]
        atnh = [at['nH'] for at in self.atoms]
        types = list(set(attype))
        nh = [-1] * len(types)
        for at, n in zip(attype, atnh):
            i = types.index(at)
            if nh[i] == -1:
                nh[i] = n
            elif nh[i] != n:
                print(f'atom type {at} with different numbers of bonded H')
                exit(1)

    def polarize(self, ff, eps=0.02):
        '''add Drude particles and bonds to topology'''

        dplist = []
        for res in ff.root.iter('Residue'):
            resname = res.get('name').replace('-', '').replace('+', '')[:3]
            for ffatom in res.iter('Atom'):
                if 'core' in ffatom.attrib:
                    dp = {'name': ffatom.get('name'),
                          'core': ffatom.get('core'),
                          'mol': resname}
                    dplist.append(dp)

        # add DP after each core
        natoms = len(self.atoms)
        random.seed(1234)
        i = 0
        self.ndrude = 0
        while i < natoms:
            atom = self.atoms[i]
            for dp in dplist:
                if atom['mol'] == dp['mol'] and atom['name'] == dp['core']:
                    ndp = {}
                    ndp['name'] = dp['name']
                    ndp['mol'] = atom['mol']
                    ndp['seq'] = atom['seq']
                    ndp['x'] = atom['x'] + eps * (2*random.random() - 1.0)
                    ndp['y'] = atom['y'] + eps * (2*random.random() - 1.0)
                    ndp['z'] = atom['z'] + eps * (2*random.random() - 1.0)
                    ndp['ele'] = 'EP'
                    i += 1
                    self.atoms.insert(i, ndp)
                    self.ndrude += 1
                    natoms += 1
                    # increment indices in bonds
                    for bond in self.bonds:
                        if bond['i'] >= i:
                            bond['i'] += 1
                        if bond['j'] >= i:
                            bond['j'] += 1
                    # no need to add DC-CP bonds
                    #self.bonds.append({'i': i - 1, 'j': i})
                    break
            i += 1

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

fpe0 =  0.0007197587     # (4 Pi eps0) in e^2/(kJ/mol A)

class Drude(object):
    """specification of drude oscillator types"""

    def __init__(self, drudefile, config, ff):
        self.types = []
        self.alpha_H = 0.323           # PCCP 20(2018)10992

        with open(drudefile, "r") as f:

            # search for polarizability of H
            for line in f:
                line = line.strip()
                if line.startswith('#') or len(line) == 0:
                    continue
                if line.startswith('H'):
                    tok = line.split()
                    self.alpha_H = float(tok[4])
                    break
            f.seek(0)

            config.countbondedH()
            config.settypes(ff)
            config.checknH()

            for line in f:
                line = line.strip()
                if line.startswith('#') or len(line) == 0:
                    continue
                tok = line.split()

                if tok[0].startswith('H'):
                    if float(tok[4]) != self.alpha_H:
                        print('warning: different polarizabilities for H')
                    continue

                drude = {}
                drude['type'] = tok[0]
                drude['dm'] = float(tok[1])
                # dq = float(tok[2])
                drude['k'] = float(tok[3])
                alpha = float(tok[4])
                drude['thole'] = float(tok[5])

                for att in config.atoms:
                    if drude['type'] == att['type']:
                        alpha += att['nH'] * self.alpha_H
                        break
                drude['alpha'] = alpha

                dq = (fpe0 * drude['k'] * alpha)**0.5
                drude['dq'] = -abs(dq)

                self.types.append(drude)


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


class Forcefield(object):
    '''force field from OpenMM xml file'''

    def __init__(self, infile):
        '''element tree from xml file'''

        self.ftree = ET.parse(infile)
        self.root = self.ftree.getroot()

        atomtypes = self.root.find('AtomTypes')
        for atom in atomtypes:
            atname = atom.get('name')
            tok = atname.split('-')
            if len(tok) < 3:
                print('error: input xml file does not have unique atom names; use fftool with --types')
                exit(1)

    def write(self, outfile):
        '''write force field to xml file'''

        for res in self.root.iter('Residue'):
            for atom in res.iter('Atom'):
                if 'core' in atom.attrib:
                    atom.attrib.pop('core')

        ET.indent(self.root, space=' ')
        self.ftree.write(outfile)

    def polarize(self, drude):
        '''add Drude dipoles to force field'''
        
        atomtypes = self.root.find('AtomTypes')
        for atom in atomtypes:
            atname = atom.get('name')
            tok = atname.split('-')
            # avoid recursion if mol name same as atom type (ex Cl)
            if atname[0] == 'D' and len(tok) > 3:
                continue
            ffname = tok[1]
            for dt in drude.types:
                if ffname == dt['type']:
                    # add DP to AtomTypes
                    dp = ET.SubElement(atomtypes, 'Type')
                    dp.set('name', 'D-' + atname)
                    dp.set('class', 'DRUD')
                    dp.set('mass', '0.0')
                    # no need to set DP, DC masses: done by OpenMM
                    # dp.set('mass', str(dt['dm']))
                    # subtract Drude mass from DC
                    # dcm = float(atom.attrib['mass'])
                    # dcm -= dt['dm']
                    # atom.set('mass', f'{dcm:.4f}')
                    break

        # check length of atom names (to decide on DP numbering)
        longname = False
        for residue in self.root.iter('Residue'):
            for atom in residue.iter('Atom'):
                if len(atom.get('name')) > 3:
                    longname = True

        for residue in self.root.iter('Residue'):
            i = 0
            cnt = 0     # counter of DP
            for atom in residue.iter('Atom'):
                i += 1
                atname = atom.get('name')
                attype = atom.get('type')
                tok = attype.split('-')
                # avoid recursion if mol name same as atom type (ex Cl)
                if attype[0] == 'D' and len(tok) > 3:
                    continue
                fftype = tok[1]
                for dt in drude.types:
                    if fftype == dt['type']:
                        cnt += 1
                        dcq = float(atom.get('charge'))
                        # subtract Drude charge from DC
                        dcq -= dt['dq']
                        atom.set('charge', f'{dcq:.5f}')
                        # insert DP
                        dp = ET.Element('Atom')
                        if not longname:
                            dp.set('name', 'D' + atname)
                        else:
                            dp.set('name', 'D' + np.base_repr(cnt, base=32))
                        dp.set('core', atname)
                        dp.set('type', 'D-' + attype)
                        dp.set('charge', '{0:.5f}'.format(dt['dq']))
                        residue.insert(i, dp)
                        # no need to add DC-DP bonds: done by OpenMM
                        #bd = ET.Element('Bond')
                        #bd.set('atomName1', atname)
                        #bd.set('atomName2', 'D' + atname)
                        #residue.append(bd)
                        break

        bdforce = self.root.find('HarmonicBondForce')
        db = ET.SubElement(bdforce, 'Bond')
        db.set('class1', 'X')
        db.set('class2', 'DRUD')
        db.set('length', '0.0')
        db.set('k', '{0:.2f}'.format(drude.types[0]['k'] * 100.0))

        nonbondedforce = self.root.find('NonbondedForce')
        dp = ET.SubElement(nonbondedforce, 'Atom')
        dp.set('class', 'DRUD')
        dp.set('sigma', '1.0')
        dp.set('epsilon', '0.0')

        drudeforce = ET.SubElement(self.root, 'DrudeForce')
        for atom in self.root.iter('Type'):
            atname = atom.get('name')
            tok = atname.split('-')
            if atname[0] == 'D' and len(tok) > 3 :
                continue
            ffname = tok[1]            
            for dt in drude.types:
                if ffname == dt['type']:
                    dp = ET.SubElement(drudeforce, 'Particle')
                    dp.set('type1', 'D-' + atname)
                    dp.set('type2', atname)
                    dp.set('charge', '{0:.5f}'.format(dt['dq']))
                    dp.set('polarizability', '{0:.7f}'.format(dt['alpha']/1e3))
                    dp.set('thole', str(dt['thole']/2))
                    break

        ljforce = self.root.find('LennardJonesForce')
        if ljforce:
            dp = ET.SubElement(ljforce, 'Atom')
            dp.set('class', 'DRUD')
            dp.set('sigma', '1.0')
            dp.set('epsilon', '0.0')


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

def main():
    parser = argparse.ArgumentParser(description = usage,
             formatter_class = argparse.RawTextHelpFormatter)
    parser.add_argument('-a', '--alpha_file', default = 'alpha.ff',
                        help = 'Drude parameter file (default: alpha.ff)')
    parser.add_argument('-e', '--eps', type=float, default = '0.02',
                        help = 'Max DC-DP distance (default: 0.02 A)')
    parser.add_argument('-ix', '--inxml', default = 'field.xml',
                        help = 'input OpenMM xml file (default: field.xml)')
    parser.add_argument('-ox', '--outxml', default = 'field-p.xml',
                        help = 'output OpenMM xml file (default: field-p.xml)')
    parser.add_argument('-ip', '--inpdb', default = 'config.pdb',
                        help = 'PDB or PDBx/mmcif file with configuration (default: config.pdb)')
    parser.add_argument('-op', '--outpdb', default = 'config-p.pdb',
                        help = 'PDB or PDBx/mmcif file with configuration (default: config-p.pdb)')
    args = parser.parse_args()

    config = Topology(args.inpdb)
    ff = Forcefield(args.inxml)
    drude = Drude(args.alpha_file, config, ff)
    ff.polarize(drude)
    config.polarize(ff, args.eps)
    print(f'{config.natom} atoms {config.ndrude} Drude particles')
    ff.write(args.outxml)       # removes 'core', so use after config.polarize()
    print('polarizable force field written to', args.outxml)

    if args.inpdb == 'config.mmcif':
        config.writepdbx(args.outpdb)
    
    if args.inpdb == 'config.pdb':
        pdb_limit = (config.natom + config.ndrude)
        if pdb_limit > 99999:
            raise RuntimeError("System too large for PDB format, use PDBx/mmcif files")
        else:
            config.writepdb(args.outpdb)

    print('configuration written to', args.outpdb)

if __name__ == '__main__':
    main()