Merge pull request #256 from rmatsum836/xml_writer

Remove undefined atom types from SMARTS strings in xml writer

foyer/tests/files/ethane-multiple.xml

@@ -0,0 +1,24 @@
+ <ForceField>
+     <!-- Forcefield xml to test _update_defs in xml_writer.py -->
+  <AtomTypes>
+    <Type name="opls_135" class="CT" element="C" mass="12.01100" def="[C;X4](C)(H)(H)H" desc="alkane CH3" doi="10.1021/ja9621760" overrides="dummy_001"/>
+    <Type name="opls_140" class="HC" element="H" mass="1.00800" def="H[%opls_135,%dummy_001]" desc="alkane H" doi="10.1021/ja9621760"/>
+    <Type name="dummy_001" class="" element="C" mass="12.01100" def="C" desc="Dummy for test_load_xml"/>
+ </AtomTypes>
+ <HarmonicBondForce>
+    <Bond class1="CT" class2="CT" length="0.1529" k="224262.4"/>
+    <Bond class1="CT" class2="HC" length="0.109" k="284512.0"/>
+ </HarmonicBondForce>
+ <HarmonicAngleForce>
+    <Angle class1="CT" class2="CT" class3="HC" angle="1.93207948196" k="313.8"/>
+    <Angle class1="HC" class2="CT" class3="HC" angle="1.88146493365" k="276.144"/>
+ </HarmonicAngleForce>
+ <RBTorsionForce>
+    <Proper class1="HC" class2="CT" class3="CT" class4="HC" c0="0.6276" c1="1.8828" c2="0.0" c3="-2.5104" c4="0.0" c5="0.0"/>
+ </RBTorsionForce>
+ <NonbondedForce coulomb14scale="0.5" lj14scale="0.5">
+     <Atom type="opls_135" charge="-0.18" sigma="0.35" epsilon="0.276144"/>
+     <Atom type="opls_140" charge="0.06" sigma="0.25" epsilon="0.12552"/>
+     <Atom type="dummy_001" charge="0.06" sigma="0.25" epsilon="0.12552"/>
+ </NonbondedForce>
+</ForceField>

foyer/tests/test_forcefield.py

@@ -376,6 +376,18 @@ def test_write_xml_multiple_periodictorsions(filename):
     assert 'k2' in periodic_element[0].attrib
     assert 'phase2' in periodic_element[0].attrib
 
+@pytest.mark.parametrize("filename", ['ethane.mol2', 'benzene.mol2'])
+def test_load_xml(filename):
+    mol = pmd.load_file(get_fn(filename), structure=True)
+    if filename == 'ethane.mol2':
+        ff = Forcefield(get_fn('ethane-multiple.xml'))
+    else:
+        ff = Forcefield(name='oplsaa')
+    typed = ff.apply(mol)
+    typed.write_foyer(filename='snippet.xml', forcefield=ff, unique=True)
+
+    generated_ff = Forcefield('snippet.xml')
+
 def test_write_xml_overrides():
     #Test xml_writer new overrides and comments features
     mol = pmd.load_file(get_fn('styrene.mol2'), structure=True)

foyer/xml_writer.py

@@ -2,6 +2,9 @@
 
 import collections
 from lxml import etree as ET
+from foyer.smarts_graph import SMARTSGraph
+import networkx as nx
+import warnings
 
 import numpy as np
 
@@ -124,6 +127,37 @@ def _write_atoms(self, root, atoms, forcefield, unique):
         nb_force.set('sigma', str(round(atom.atom_type.sigma/10, 4)))
         nb_force.set('epsilon', str(round(atom.atom_type.epsilon * 4.184, 6)))
 
+    _update_defs(atomtypes, nonbonded, forcefield)
+
+def _update_defs(atomtypes, nonbonded, forcefield):
+    def_list = [i.get('def') for i in atomtypes.iterchildren()]
+    name_list = [i.get('name') for i in atomtypes.iterchildren()]
+    smarts_list = list()
+    smarts_parser = forcefield.parser
+    for smarts_string, name in zip(def_list, name_list):
+        smarts_graph = SMARTSGraph(smarts_string, parser=smarts_parser,
+                                   name=name)
+        for atom_expr in nx.get_node_attributes(smarts_graph, name='atom').values():
+            labels = atom_expr.find_data('has_label')
+            for label in labels:
+                atom_type = label.children[0][1:]
+                smarts_list.append(atom_type)
+    smarts_list = list(set(smarts_list))
+    extra_types = [i for i in smarts_list if i not in name_list]
+
+    for extra in extra_types:
+        for i, definition in enumerate(def_list):
+            if extra in definition:
+                warnings.warn('Removing undefined atom type `{}`'
+                        ' from SMARTS string `{}`'.format(
+                            extra, definition))
+                extra_edit = '%' + extra
+                extra_index = definition.find(extra_edit)
+                if definition[extra_index-1] == ';':
+                    new_def = definition.replace(extra_edit + ',' , '')
+                else:
+                    new_def = definition.replace(',' + extra_edit, '')
+                atomtypes[i].set('def', new_def)
 
 def _write_bonds(root, bonds, unique):
     bond_forces = ET.SubElement(root, 'HarmonicBondForce')