Force Field Fitting based on DMFF- I. Topology Transformation

引言

​ 最近由于科研需要，开始接触力场拟合。DMFF(Differentiable Molecular Force Field)是DeepModeling开发，基于JAX自动微分技术对分子力场自动调参的平台。其理念对厌恶繁琐事物的我来说十分有吸引力（本人有将一切事情都自动化、黑箱化的倾向）。

​ DMFF中力和能量的计算是通过对OpenMM中ForceField类的重载来实现的，在DMFF中就是Hamiltonian类。为了计算一个结构的力和能量，需要读取键连关系(loadBondDefinitions)、读取力场以及目标结构原子类型的定义(Hamiltonian)、并根据原子位置创建计算势能和力(createPotential)，其中键连关系和力场、残基原子类型定义都是通过OpenMM格式的xml文件来储存。

​ 参考OpenMM官方文档、一个gmx2xml脚本和ParmEd的功能，写了转换脚本prmtop2xml.py，并用test_trans.py检验正确性。注意OpenMM读取转换得到的xml文件计算的非键相互作用能与DMFF、ParmEd的不同是因为prmtop中相同原子类型的原子也可以有不同的电荷，但在OpenMM forcefield中一种原子类型只能有一种电荷，DMFF可以从残基定义的xml文件中单独读取每个原子的电荷，但OpenMM不能，除非为每一个原子定义一个原子类型。

转换脚本和测试脚本

Creating Force Fields in OpenMM

OpenMM模拟蛋白-配体体系

ParmEd Amber files

prmtop2xml.py

注意，要使得OpenMM或DMFF正确读取二面角项，必须要这样写：

import parmed as chem
import parmed.unit as u
import openmm.app as app
import openmm as mm
from sys import stdout
import xml.etree.ElementTree as ET
from datetime import datetime
from xml.etree.ElementTree import Element
from xml.etree.ElementTree import SubElement
from xml.etree.ElementTree import ElementTree
from math import pi

# elemnt为传进来的Elment类，参数indent用于缩进，newline用于换行   
def prettyXml(element, indent, newline, level = 0): 
    # 判断element是否有子元素
    if element:
        # 如果element的text没有内容      
        if element.text == None or element.text.isspace():     
            element.text = newline + indent * (level + 1)      
        else:    
            element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * (level + 1)    
    # 此处两行如果把注释去掉，Element的text也会另起一行 
    #else:     
        #element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * level    
    temp = list(element) # 将elemnt转成list    
    for subelement in temp:    
        # 如果不是list的最后一个元素，说明下一个行是同级别元素的起始，缩进应一致
        if temp.index(subelement) < (len(temp) - 1):     
            subelement.tail = newline + indent * (level + 1)    
        else:  # 如果是list的最后一个元素， 说明下一行是母元素的结束，缩进应该少一个    
            subelement.tail = newline + indent * level   
        # 对子元素进行递归操作 
        prettyXml(subelement, indent, newline, level = level + 1)     
        

def get_element(atomics_number):
    # Convert atomic number to element symbol
    element_dict = {
        1: 'H', 6: 'C', 7: 'N', 8: 'O', 9: 'F',
        15: 'P', 16: 'S', 17: 'Cl', 35: 'Br', 53: 'I'
    }
    return element_dict.get(atomics_number, None)

def write_prm_xml(xml_name):
    #Basic Information of Residues XML
    ForceField = Element('ForceField')
    info = SubElement(ForceField, 'Info')
    DataGenerated=SubElement(info,"DataGenerated")
    DataGenerated.text=str(datetime.today())
    Reference=SubElement(info,"Reference")
    Reference.text="prmtop2xml was developed by csming; 2025.04.20"


    # Residues
    Residues=SubElement(ForceField,"Residues")
    for residue in parm.residues:
        Residue=SubElement(Residues,"Residue")
        Residue.set("name",residue.name)
        for atom in residue.atoms:
            Atom=SubElement(Residue,"Atom")
            Atom.set("name",atom.name)
            Atom.set("type",atom.type)
            Atom.set("charge",str(atom.charge))
    for bond in parm.bonds:
        if bond.atom1.residue.idx==bond.atom2.residue.idx:
            Bond=SubElement(Residues.find("./Residue[@name='{}']"\
                                          .format(bond.atom1.residue.name)),
                                          "Bond")
            Bond.set("atomName1",bond.atom1.name)
            Bond.set("atomName2",bond.atom2.name)
        else:
            ExternalBond1=SubElement(Residues.find("./Residue[@name='{}']"\
                                          .format(bond.atom1.residue.name)),
                                          "ExternalBond")
            ExternalBond1.set("atomName",bond.atom1.name)
            ExternalBond2=SubElement(Residues.find("./Residue[@name='{}']"\
                                          .format(bond.atom2.residue.name)),
                                          "ExternalBond")
            ExternalBond2.set("atomName",bond.atom2.name)

    tree = ElementTree(ForceField)
    root = tree.getroot()                  #得到根元素，Element类    
    prettyXml(root, '\t', '\n')            #执行美化方法    


    with open(xml_name,'w') as xml_file:
        # write out xml data
        tree.write(xml_name, encoding = 'utf-8')
    print("prmtop file has been converted into xml file")


def write_top_xml(xml_name):

    # Residues
    Residues=Element("Residues")
    for residue in parm.residues:
        Residue=SubElement(Residues,"Residue")
        Residue.set("name",residue.name)
    for bond in parm.bonds:
        if bond.atom1.residue.idx==bond.atom2.residue.idx:
            Bond=SubElement(Residues.find("./Residue[@name='{}']"\
                                          .format(bond.atom1.residue.name)),
                                          "Bond")
            Bond.set("from",bond.atom1.name)
            Bond.set("to",bond.atom2.name)

    tree = ElementTree(Residues)
    root = tree.getroot()                  #得到根元素，Element类    
    prettyXml(root, '\t', '\n')            #执行美化方法   


    with open(xml_name,'w') as xml_file:
        # write out xml data
        tree.write(xml_name, encoding = 'utf-8')
    print("prmtop file has been converted into xml file")

def write_forcefield_xml(xml_name,external_charge=False):
    #Basic Information of Residues XML
    ForceField = Element('ForceField')
    info = SubElement(ForceField, 'Info')
    DataGenerated=SubElement(info,"DataGenerated")
    DataGenerated.text=str(datetime.today())
    Reference=SubElement(info,"Reference")
    Reference.text="prmtop2xml was developed by csming; 2025.04.20"

    #Atomtypes
    Atomtypes=SubElement(ForceField,"AtomTypes")
    atomtypes_dict=dict(zip([atom.type for atom in parm.atoms],\
                            [(atom.element_name,atom.mass,\
                              atom.charge,atom.sigma,atom.epsilon) for atom in parm.atoms]))
    atomtypes_dict_sorted=dict(sorted(atomtypes_dict.items(),key=lambda \
s:s[0]))
    for atomtype,(atom_element,mass,charge,sigma,epsilon) in atomtypes_dict_sorted.items():
        Type=SubElement(Atomtypes,"Type")   
        Type.set("element",atom_element)
        Type.set("name",atomtype)
        Type.set("class",atomtype)

        Type.set("mass",str(mass))
    '''
    # Residues
    Residues=SubElement(ForceField,"Residues")
    for residue in parm.residues:
        Residue=SubElement(Residues,"Residue")
        Residue.set("name",residue.name)
        for atom in residue.atoms:
            Atom=SubElement(Residue,"Atom")
            Atom.set("name",atom.name)
            Atom.set("type",atom.type)
            Atom.set("charge",str(atom.charge))
    for bond in parm.bonds:
        if bond.atom1.residue.idx==bond.atom2.residue.idx:
            Bond=SubElement(Residues.find("./Residue[@name='{}']"\
                                          .format(bond.atom1.residue.name)),
                                          "Bond")
            Bond.set("atomName1",bond.atom1.name)
            Bond.set("atomName2",bond.atom2.name)
        else:
            ExternalBond1=SubElement(Residues.find("./Residue[@name='{}']"\
                                          .format(bond.atom1.residue.name)),
                                          "ExternalBond")
            ExternalBond1.set("atomName",bond.atom1.name)
            ExternalBond2=SubElement(Residues.find("./Residue[@name='{}']"\
                                          .format(bond.atom2.residue.name)),
                                          "ExternalBond")
            ExternalBond2.set("atomName",bond.atom2.name)
    '''
    #HarmonicBondForce
    bonds_dict=dict(zip([(bond.atom1.type,bond.atom2.type) for bond in parm.bonds],\
                         [(bond.type.k,bond.type.req) for bond in parm.bonds]))
    new_bonds_dict = bonds_dict.copy()
    #for key in bonds_dict.keys():
    #    if (key[1],key[0]) in new_bonds_dict.keys():
    #        del new_bonds_dict[key]
    
    HarmonicBondForce=SubElement(ForceField,"HarmonicBondForce")
    for bond,(k,req) in new_bonds_dict.items():
        Bond=SubElement(HarmonicBondForce,"Bond")
        Bond.set("class1",bond[0])
        Bond.set("class2",bond[1])
        #parmed: when red from from prmtop, k(r-req)^2 
        # Force constant in kcal/mol/Angstrom^2 
        # Equilibrium bond distance in Angstroms

        #openmm: 1/2 k (r-req)^2 
        # Force constant in kJ/mol/nm^2
        # Equilibrium bond distance in nm
        #1 kcal/mol = 4.184 kJ
        #1 Angstrom = 0.1 nm
        Bond.set("length",str(req*0.1))
        Bond.set("k",str(k*4.184*100*2))

    
    #HarmonicAngleForce
    angles_dict=dict(zip([(angle.atom1.type,angle.atom2.type,angle.atom3.type) for angle in parm.angles],\
                         [(angle.type.k,angle.type.theteq) for angle in parm.angles]))
    new_angles_dict = angles_dict.copy()
    #for key in angles_dict.keys():
    #    if (key[2],key[1],key[0]) in new_angles_dict.keys():
    #        del new_angles_dict[key]
    HarmonicAngleForce=SubElement(ForceField,"HarmonicAngleForce")
    for angle,(k,theteq) in new_angles_dict.items():
        Angle=SubElement(HarmonicAngleForce,"Angle")
        Angle.set("class1",angle[0])
        Angle.set("class2",angle[1])
        Angle.set("class3",angle[2])
        #parmed: when red from from prmtop, k(theta-theteq)^2 
        # Force constant in kcal/mol/radian^2 
        # Equilibrium angle in degrees

        #openmm: 1/2 k (theta-theteq)^2 
        # Force constant in kJ/mol/radian^2
        # Equilibrium angle in radians
        #1 kcal/mol = 4.184 kJ
        #1 degree = pi/180 radian
        Angle.set("angle",str(theteq*(pi/180)))
        Angle.set("k",str(k*4.184*2))

    #PeriodicTorsionForce

    dihedrals_dict = {
        tuple((dihedral.improper,
               dihedral.atom1.type,
               dihedral.atom2.type,
               dihedral.atom3.type,
               dihedral.atom4.type)): []
        for dihedral in parm.dihedrals
    }

    for dihedral in parm.dihedrals:
        key = [dihedral.improper,\
               dihedral.atom1.type,dihedral.atom2.type,\
              dihedral.atom3.type,dihedral.atom4.type]
        value=[dihedral.type.phi_k,dihedral.type.phase,dihedral.type.per]
        
        if value not in dihedrals_dict[tuple(key)]:
            dihedrals_dict[tuple(key)].append(value)
            
    new_dihedrals_dict = dihedrals_dict.copy()

    #for key in dihedrals_dict.keys():
    #    if key[0] == 'improper':
    #        #improper dihedrals
    #        permutation = [(key[0],key[2],key[1],key[3],key[4]),\
    #                       (key[0],key[4],key[1],key[3],key[2]),\
    #                        (key[0],key[4],key[2],key[3],key[1]),\
    #                        (key[0],key[2],key[4],key[3],key[1]),\
    #                        (key[0],key[1],key[4],key[3],key[2])]
    #        if any(permutation[i] in new_dihedrals_dict.keys() for i in range(5)):
    #            del new_dihedrals_dict[key]
    #    else:  
    #        if (key[0],key[4],key[3],key[2],key[1]) in new_dihedrals_dict.keys():
    #            del new_dihedrals_dict[key]

    PeriodicTorsionForce=SubElement(ForceField,"PeriodicTorsionForce")
    PeriodicTorsionForce.set("ordering","amber")
    for dihedral,dihedral_types in new_dihedrals_dict.items():
        if dihedral[0] == True:
            #improper dihedrals
            Dihedral=SubElement(PeriodicTorsionForce,"Improper")
            Dihedral.set("type1",dihedral[1])
            Dihedral.set("type2",dihedral[2])
            Dihedral.set("type3",dihedral[3])
            Dihedral.set("type4",dihedral[4])
        else:
            #proper dihedrals
            Dihedral=SubElement(PeriodicTorsionForce,"Proper")
            Dihedral.set("type1",dihedral[1])
            Dihedral.set("type2",dihedral[2])
            Dihedral.set("type3",dihedral[3])
            Dihedral.set("type4",dihedral[4])
        
        #parmed: when read from from prmtop, phi_k [1+cos(per*phi-phase)]
        # Force constant in kcal/mol
        # Phase angle in degrees
        # improper dihedrals: an Amber-style improper torsion,
        #  where atom3 is the “central” atom bonded to atoms 1, 2, and 4 

        #openmm: k (1+cos(n*phi-phi0)) 
        # Force constant in kJ/mol
        # Phase angle in radians
        # inproper dihedrals: an OpenMM-style improper torsion,
        #  where atom1 is the “central” atom bonded to atoms 2, 3, and 4

        #1 kcal/mol = 4.184 kJ
        #1 degree = pi/180 radian
        for i,(k,phi0,per) in enumerate(dihedral_types):
            Dihedral.set("periodicity{}".format(i+1),str(per))
            Dihedral.set("phase{}".format(i+1),str(phi0*(pi/180)))
            Dihedral.set("k{}".format(i+1),str(k*4.184))
    
    #NonbondedForce
    NonbondedForce=SubElement(ForceField,"NonbondedForce")
    NonbondedForce.set("coulomb14scale","0.833333")
    NonbondedForce.set("lj14scale","0.5")
    if external_charge==True:
        UseAttributeFromResidue=SubElement(NonbondedForce,"UseAttributeFromResidue")
        UseAttributeFromResidue.set("name","charge")
    for atomtype,(atom_element,mass,charge,sigma,epsilon) in atomtypes_dict_sorted.items():
        NB_Atom=SubElement(NonbondedForce,"Atom")   
        NB_Atom.set("type",atomtype)
        NB_Atom.set("charge",str(charge))
        NB_Atom.set("sigma",str(sigma*0.1))
        NB_Atom.set("epsilon",str(epsilon*4.184))

    tree = ElementTree(ForceField)
    root = tree.getroot()                  #得到根元素，Element类  
      
    prettyXml(root, '\t', '\n')            #执行美化方法    


    with open(xml_name,'w') as xml_file:
        # write out xml data
        tree.write(xml_name, encoding = 'utf-8')
    print("prmtop file has been converted into xml file")


#parm = chem.amber.AmberParm('CHA.prmtop')
parm = chem.load_file('CHA.prmtop')

write_prm_xml('CHA-prm.xml')
write_top_xml('CHA-top.xml')
write_forcefield_xml('CHA-force_file_dmff.xml',external_charge=True)
write_forcefield_xml('CHA-force_file.xml')

test_trans.py

import jax
import jax.numpy as jnp
import openmm.app as app
import openmm as mm
import openmm.unit as unit
import parmed as chem
from dmff import Hamiltonian, NeighborList

app.Topology.loadBondDefinitions("CHA-top.xml")
pdb = app.PDBFile("CHA.pdb")
#ff = Hamiltonian("gaff-2.11.xml","lig-prm.xml")
ff = Hamiltonian("CHA-force_file_dmff.xml","CHA-prm.xml")
potentials = ff.createPotential(pdb.topology)

params = ff.getParameters()
nbparam = params['NonbondedForce']
nbparam

positions = jnp.array(pdb.getPositions(asNumpy=True).value_in_unit(unit.nanometer))
box = jnp.array([
    [10.0,  0.0,  0.0],
    [ 0.0, 10.0,  0.0],
    [ 0.0,  0.0, 10.0]
])
nbList = NeighborList(box, rcut=4,cov_map=potentials.meta["cov_map"])
nbList.allocate(positions)
pairs = nbList.pairs

for component in potentials.dmff_potentials:
    func=potentials.dmff_potentials[component]
    ene = func(positions, box, pairs, params)
    print(component, ene)

efunc = potentials.getPotentialFunc()
params = ff.getParameters()
totene = efunc(positions, box, pairs, params)

print("Total Potential Energy",totene)

#############################################
print("\n")
print(("Energy Calculate From OpenMM"))

app.Topology.loadBondDefinitions("CHA-top.xml")
pdb = app.PDBFile("CHA.pdb")
ff2 = app.ForceField("CHA-force_file.xml","CHA-prm.xml")

system2 = ff2.createSystem(pdb.topology,nonbondedMethod=app.NoCutoff)

for i, f in enumerate(system2.getForces()):
    f.setForceGroup(i)

integrator2 = mm.LangevinMiddleIntegrator(300*unit.kelvin, 1/unit.picosecond, 0.004*unit.picoseconds)
simulation2 = app.Simulation(pdb.topology, system2, integrator2)
simulation2.context.setPositions(pdb.positions)


for i, f in enumerate(system2.getForces()):
    state = simulation2.context.getState(getEnergy=True, groups={i})
    print(f.getName(), state.getPotentialEnergy())

state = simulation2.context.getState(energy=True)
simulation2.context.getState(energy=True)
print("Total Potential Energy",state.getPotentialEnergy())

########################################################

pdb = app.PDBFile('CHA.pdb')
prmtop = app.AmberPrmtopFile('CHA.prmtop')
parm = chem.load_file('CHA.prmtop')

system0 = prmtop.createSystem(nonbondedMethod=app.NoCutoff)

for i, f in enumerate(system0.getForces()):
    f.setForceGroup(i)

integrator = mm.LangevinMiddleIntegrator(300*unit.kelvin, 1/unit.picosecond, 0.004*unit.picoseconds)
simulation0 = app.Simulation(prmtop.topology, system0, integrator)
simulation0.context.setPositions(pdb.positions)
#print("\n")
#print("Energy Decomposition with parm")
#
#struct0 = chem.openmm.load_topology(pdb.topology, system0)
#for item in chem.openmm.energy_decomposition(struct0, simulation0.context).items():
#    print(item[0], item[1]*4.184)

print("\n")

print(("Energy Calculate From prmtop"))

for i, f in enumerate(system0.getForces()):
    state = simulation0.context.getState(getEnergy=True, groups={i})
    print(f.getName(), state.getPotentialEnergy())

state = simulation0.context.getState(energy=True)
print("Total Potential Energy",state.getPotentialEnergy())

########################################################

pdb = app.PDBFile('CHA.pdb')
parm = chem.load_file('CHA.prmtop')

system0 = parm.createSystem(nonbondedMethod=app.NoCutoff)


integrator = mm.LangevinMiddleIntegrator(300*unit.kelvin, 1/unit.picosecond, 0.004*unit.picoseconds)
simulation0 = app.Simulation(parm.topology, system0, integrator)
simulation0.context.setPositions(pdb.positions)
print("\n")
print("Energy Decomposition with parm")

struct0 = chem.openmm.load_topology(pdb.topology, system0)
for item in chem.openmm.energy_decomposition(struct0, simulation0.context).items():
    print(item[0], item[1]*4.184)

bonded = [f for f in system0.getForces() if isinstance(f, mm.PeriodicTorsionForce)][0]
for i in range(bonded.getNumTorsions()):
    particle1, particle2, particle3, particle4, per, length, k = bonded.getTorsionParameters(i)
    if 1==1:#particle1 == 0 or particle2 == 0 or particle3 == 0 or particle4 == 0:
        print(f'Particles ({particle1}, {particle2},{particle2},{particle4}), periodicity = {per}, length = {length}, k = {k}')
print(bonded.getNumTorsions())
print("openmm")

bonded = [f for f in system2.getForces() if isinstance(f, mm.PeriodicTorsionForce)][0]
for i in range(bonded.getNumTorsions()):
    particle1, particle2, particle3, particle4, per, length, k = bonded.getTorsionParameters(i)
    if 1==1:#particle1 == 0 or particle2 == 0 or particle3 == 0 or particle4 == 0:
        print(f'Particles ({particle1}, {particle2},{particle2},{particle4}), periodicity = {per}, length = {length}, k = {k}')

print(bonded.getNumTorsions()