diff --git a/dmff/api.py b/dmff/api.py index d9bec951b..ba9823ff3 100644 --- a/dmff/api.py +++ b/dmff/api.py @@ -6,6 +6,7 @@ import numpy as np import jax.numpy as jnp from collections import defaultdict +import xml.etree.ElementTree as ET from .admp.disp_pme import ADMPDispPmeForce from .admp.multipole import convert_cart2harm, rot_local2global from .admp.pairwise import TT_damping_qq_c6_kernel, generate_pairwise_interaction @@ -31,6 +32,39 @@ import sys +class XMLNodeInfo: + + class XMLElementInfo: + + def __init__(self, name): + self.name = name + self.attributes = {} + + def addAttribute(self, key, value): + self.attributes[key] = value + + + def __init__(self, name): + self.name = name + self.attributes = {} + self.elements = [] + + + def addAttribute(self, key, value): + self.attributes[key] = value + + + def addElement(self, name, info): + element = self.XMLElementInfo(name) + for k, v in info.items(): + element.addAttribute(k, v) + self.elements.append(element) + + + def modResidue(self, residue, atom, key, value): + pass + + def get_line_context(file_path, line_number): return linecache.getline(file_path, line_number).strip() @@ -56,6 +90,15 @@ def build_covalent_map(data, max_neighbor): return covalent_map +def findAtomTypeTexts(attribs, num): + typetxt = [] + for n in range(1, num+1): + for key in ["type%i"%n, "class%i"%n]: + if key in attribs: + typetxt.append((key, attribs[key])) + break + return typetxt + class ADMPDispGenerator: def __init__(self, hamiltonian): self.ff = hamiltonian @@ -630,10 +673,13 @@ def __init__(self, hamiltonian): self.params = {"k": [], "length": []} self._jaxPotential = None self.types = [] + self.typetexts = [] def registerBondType(self, bond): + typetxt = findAtomTypeTexts(bond, 2) types = self.ff._findAtomTypes(bond, 2) self.types.append(types) + self.typetexts.append(typetxt) self.params["k"].append(float(bond["k"])) self.params["length"].append(float(bond["length"])) # length := r0 @@ -647,10 +693,9 @@ def parseElement(element, hamiltonian): - <\HarminicBondForce> + <\HarmonicBondForce> """ - generator = HarmonicBondJaxGenerator(hamiltonian) hamiltonian.registerGenerator(generator) for bondtype in element.findall("Bond"): @@ -700,7 +745,18 @@ def getJaxPotential(self): def renderXML(self): # generate xml force field file - pass + finfo = XMLNodeInfo("HarmonicBondForce") + for ntype in range(len(self.types)): + binfo = {} + k1, v1 = self.typetexts[ntype][0] + k2, v2 = self.typetexts[ntype][1] + binfo[k1] = v1 + binfo[k2] = v2 + for key in self.params.keys(): + binfo[key] = "%.8f"%self.params[key][ntype] + finfo.addElement("Bond", binfo) + return finfo + # register all parsers @@ -1450,7 +1506,10 @@ def createForce(self, system, data, nonbondedMethod, nonbondedCutoff, args): self.params[k] = jnp.array(self.params[k]) mscales_coul = jnp.array([0.0, 0.0, 0.0, 1.0, 1.0, 1.0]) # mscale for PME - mscales_coul = mscales_coul.at[2].set(1.0 - self.params["coulomb14scale"][0]) + mscales_coul = mscales_coul.at[2].set(self.params["coulomb14scale"][0]) + mscales_lj = jnp.array([0.0, 0.0, 0.0, 1.0, 1.0, 1.0]) # mscale for LJ + mscales_lj = mscales_lj.at[2].set(self.params["lj14scale"][0]) + # Coulomb: only support PME for now # set PBC @@ -1505,29 +1564,6 @@ def createForce(self, system, data, nonbondedMethod, nonbondedCutoff, args): map_nbfix = np.array(map_nbfix, dtype=int).reshape((-1, 2)) colv_map = build_covalent_map(data, 6) - map_exclusion = [] - scale_14 = [] - npair = 0 - for ii in range(colv_map.shape[0]): - for jj in range(ii + 1, colv_map.shape[1]): - if colv_map[ii, jj] > 0 and colv_map[ii, jj] < 4: - map_exclusion.append((ii, jj)) - if colv_map[ii, jj] == 3: - scale_14.append(npair) - npair += 1 - - map_exclusion = np.array(map_exclusion, dtype=int) - scale_14 = np.array(scale_14, dtype=int) - scale_lj_exclusion = np.ones((map_exclusion.shape[0],)) - scale_coul_exclusion = np.ones((map_exclusion.shape[0],)) - scale_lj_exclusion = jnp.array(scale_lj_exclusion) - scale_coul_exclusion = jnp.array(scale_coul_exclusion) - scale_lj_exclusion = scale_lj_exclusion.at[scale_14].set( - 1.0 - self.params["lj14scale"][0] - ) - scale_coul_exclusion = scale_coul_exclusion.at[scale_14].set( - 1.0 - self.params["coulomb14scale"][0] - ) if unit.is_quantity(nonbondedCutoff): r_cut = nonbondedCutoff.value_in_unit(unit.nanometer) @@ -1549,8 +1585,7 @@ def createForce(self, system, data, nonbondedMethod, nonbondedCutoff, args): r_cut, map_lj, map_nbfix, - map_exclusion, - scale_lj_exclusion, + colv_map, isSwitch=ifSwitch, isPBC=ifPBC, isNoCut=isNoCut, @@ -1561,14 +1596,10 @@ def createForce(self, system, data, nonbondedMethod, nonbondedCutoff, args): # do not use PME if nonbondedMethod in [app.CutoffPeriodic, app.CutoffNonPeriodic]: # use Reaction Field - coulforce = CoulReactionFieldForce( - r_cut, map_charge, map_exclusion, scale_coul_exclusion, isPBC=ifPBC - ) + coulforce = CoulReactionFieldForce(r_cut, map_charge, colv_map, isPBC=ifPBC) if nonbondedMethod is app.NoCutoff: # use NoCutoff - coulforce = CoulNoCutoffForce( - map_charge, map_exclusion, scale_coul_exclusion - ) + coulforce = CoulNoCutoffForce(map_charge, colv_map) else: coulforce = CoulombPMEForce(box, r_cut, self.ethresh, colv_map) @@ -1584,6 +1615,7 @@ def potential_fn(positions, box, pairs, params): params["sigma"], params["epsfix"], params["sigfix"], + mscales_lj ) coulE = coulenergy(positions, box, pairs, params["charge"], mscales_coul) @@ -1604,6 +1636,9 @@ def renderXML(self): class Hamiltonian(app.forcefield.ForceField): def __init__(self, *xmlnames): super().__init__(*xmlnames) + # add a function to parse AtomTypes and Residues information + self._atomtypes = None + self._residues = None self._potentials = [] def createPotential( @@ -1629,6 +1664,23 @@ def createPotential( pass return [p for p in self._potentials] + def render(self, filename): + root = ET.Element("ForceField") + forceInfos = [g.renderXML() for g in self._forces] + for finfo in forceInfos: + # create xml nodes + if finfo is not None: + node = ET.SubElement(root, finfo.name) + for key in finfo.attributes.keys(): + node.set(key, finfo.attributes[key]) + for elem in finfo.elements: + subnode = ET.SubElement(node, elem.name) + for key in elem.attributes.keys(): + subnode.set(key, elem.attributes[key]) + + tree = ET.ElementTree(root) + tree.write(filename) + if __name__ == "__main__": H = Hamiltonian("forcefield.xml") diff --git a/dmff/classical/inter.py b/dmff/classical/inter.py index a6ba22812..7051c3903 100644 --- a/dmff/classical/inter.py +++ b/dmff/classical/inter.py @@ -19,8 +19,7 @@ def __init__( r_cut, map_prm, map_nbfix, - map_exclusion, - scale_exclusion, + colvmap, isSwitch=False, isPBC=True, isNoCut=False, @@ -31,10 +30,9 @@ def __init__( self.map_prm = map_prm self.map_nbfix = map_nbfix - self.map_exclusion = map_exclusion - self.scale_exclusion = scale_exclusion self.ifPBC = isPBC self.ifNoCut = isNoCut + self.colvmap = colvmap def generate_get_energy(self): def get_LJ_energy(dr_vec, sig, eps, box): @@ -60,7 +58,7 @@ def get_LJ_energy(dr_vec, sig, eps, box): return E - def get_energy(positions, box, pairs, epsilon, sigma, epsfix, sigfix): + def get_energy(positions, box, pairs, epsilon, sigma, epsfix, sigfix, mscales): eps_m1 = jnp.repeat(epsilon.reshape((-1, 1)), epsilon.shape[0], axis=1) eps_m2 = eps_m1.T @@ -74,29 +72,21 @@ def get_energy(positions, box, pairs, epsilon, sigma, epsfix, sigfix): sig_mat = sig_mat.at[self.map_nbfix[:, 0], self.map_nbfix[:, 1]].set(sigfix) sig_mat = sig_mat.at[self.map_nbfix[:, 1], self.map_nbfix[:, 0]].set(sigfix) + colv_pair = self.colvmap[pairs[:,0],pairs[:,1]] + mscale_pair = mscales[colv_pair-1] # in mscale vector, the 0th item is 1-2 scale, the 1st item is 1-3 scale, etc... + + dr_vec = positions[pairs[:, 0]] - positions[pairs[:, 1]] prm_pair0 = self.map_prm[pairs[:, 0]] prm_pair1 = self.map_prm[pairs[:, 1]] eps = eps_mat[prm_pair0, prm_pair1] sig = sig_mat[prm_pair0, prm_pair1] - E_inter = get_LJ_energy(dr_vec, sig, eps, box) - - # exclusion - dr_excl_vec = ( - positions[self.map_exclusion[:, 0]] - - positions[self.map_exclusion[:, 1]] - ) - excl_map0 = self.map_prm[self.map_exclusion[:, 0]] - excl_map1 = self.map_prm[self.map_exclusion[:, 1]] - eps_excl = eps_mat[excl_map0, excl_map1] - sig_excl = sig_mat[excl_map0, excl_map1] + eps_scale = eps * mscale_pair - E_excl = get_LJ_energy(dr_excl_vec, sig_excl, eps_excl, box) - E_excl = self.scale_exclusion * E_excl + E_inter = get_LJ_energy(dr_vec, sig, eps_scale, box) - return jnp.sum(E_inter) - jnp.sum(E_excl) - #return jnp.sum(E_inter) + return jnp.sum(E_inter) return get_energy @@ -104,12 +94,11 @@ def get_energy(positions, box, pairs, epsilon, sigma, epsfix, sigfix): class CoulNoCutoffForce: # E=\frac{{q}_{1}{q}_{2}}{4\pi\epsilon_0\epsilon_1 r} - def __init__(self, map_prm, map_exclusion, scale_exclusion, epsilon_1=1.0) -> None: + def __init__(self, map_prm, colvmap, epsilon_1=1.0) -> None: self.eps_1 = epsilon_1 self.map_prm = map_prm - self.map_exclusion = map_exclusion - self.scale_exclusion = scale_exclusion + self.colvmap = colvmap def generate_get_energy(self): def get_coul_energy(dr_vec, chrgprod, box): @@ -121,30 +110,21 @@ def get_coul_energy(dr_vec, chrgprod, box): return E def get_energy(positions, box, pairs, charges, mscales): + + colv_pair = self.colvmap[pairs[:,0],pairs[:,1]] + mscale_pair = mscales[colv_pair-1] + chrg_map0 = self.map_prm[pairs[:, 0]] chrg_map1 = self.map_prm[pairs[:, 1]] charge0 = charges[chrg_map0] charge1 = charges[chrg_map1] chrgprod = charge0 * charge1 + chrgprod_scale = chrgprod * mscale_pair dr_vec = positions[pairs[:, 0]] - positions[pairs[:, 1]] - E_inter = get_coul_energy(dr_vec, chrgprod, box) - - # exclusion - dr_excl_vec = ( - positions[self.map_exclusion[:, 0]] - - positions[self.map_exclusion[:, 1]] - ) - excl_map0 = self.map_prm[self.map_exclusion[:, 0]] - excl_map1 = self.map_prm[self.map_exclusion[:, 1]] - chrg0_excl = charges[excl_map0] - chrg1_excl = charges[excl_map1] - chrgprod_excl = chrg0_excl * chrg1_excl + E_inter = get_coul_energy(dr_vec, chrgprod_scale, box) - E_excl = get_coul_energy(dr_excl_vec, chrgprod_excl, box) - E_excl = self.scale_exclusion * E_excl - - return jnp.sum(E_inter) - jnp.sum(E_excl) + return jnp.sum(E_inter) return get_energy @@ -155,8 +135,7 @@ def __init__( self, r_cut, map_prm, - map_exclusion, - scale_exclusion, + colvmap, epsilon_1=1.0, epsilon_solv=78.5, isPBC=True, @@ -168,8 +147,7 @@ def __init__( self.exp_solv = epsilon_solv self.eps_1 = epsilon_1 self.map_prm = map_prm - self.map_exclusion = map_exclusion - self.scale_exclusion = scale_exclusion + self.colvmap = colvmap self.ifPBC = isPBC def generate_get_energy(self): @@ -191,30 +169,21 @@ def get_rf_energy(dr_vec, chrgprod, box): return E def get_energy(positions, box, pairs, charges, mscales): + + colv_pair = self.colvmap[pairs[:,0],pairs[:,1]] + mscale_pair = mscales[colv_pair-1] + chrg_map0 = self.map_prm[pairs[:, 0]] chrg_map1 = self.map_prm[pairs[:, 1]] charge0 = charges[chrg_map0] charge1 = charges[chrg_map1] chrgprod = charge0 * charge1 + chrgprod_scale = chrgprod * mscale_pair dr_vec = positions[pairs[:, 0]] - positions[pairs[:, 1]] - E_inter = get_rf_energy(dr_vec, chrgprod, box) - - # exclusion - dr_excl_vec = ( - positions[self.map_exclusion[:, 0]] - - positions[self.map_exclusion[:, 1]] - ) - excl_map0 = self.map_prm[self.map_exclusion[:, 0]] - excl_map1 = self.map_prm[self.map_exclusion[:, 1]] - chrg0_excl = charges[excl_map0] - chrg1_excl = charges[excl_map1] - chrgprod_excl = chrg0_excl * chrg1_excl - - E_excl = get_rf_energy(dr_excl_vec, chrgprod_excl, box) - E_excl = self.scale_exclusion * E_excl + E_inter = get_rf_energy(dr_vec, chrgprod_scale, box) - return jnp.sum(E_inter) - jnp.sum(E_excl) + return jnp.sum(E_inter) return get_energy diff --git a/dmff/settings.py b/dmff/settings.py index 9afabd551..0ba35abe1 100644 --- a/dmff/settings.py +++ b/dmff/settings.py @@ -1,6 +1,6 @@ from jax.config import config -PRECISION = 'double' # 'double' +PRECISION = 'float' # 'double' DO_JIT = True