Skip to content
Snippets Groups Projects
Commit 569e2840 authored by James Graham's avatar James Graham
Browse files

Code cleaning in ForceField

parent 9ddf1f99
No related branches found
No related tags found
No related merge requests found
...@@ -46,7 +46,11 @@ class ForceField: ...@@ -46,7 +46,11 @@ class ForceField:
with open(os.path.join(self.dirname, "forcefield.doc"), "w") as doc: with open(os.path.join(self.dirname, "forcefield.doc"), "w") as doc:
print("PyCGTOOL produced MARTINI force field - {0}".format(name), file=doc) print("PyCGTOOL produced MARTINI force field - {0}".format(name), file=doc)
def write_rtp(self, filename, mapping, bonds): def write(self, filename, mapping, bonds):
nterms, cterms, bothterms = self._write_rtp(filename, mapping, bonds)
self._write_r2b(filename, nterms, cterms, bothterms)
def _write_rtp(self, filename, mapping, bonds):
""" """
Write a GROMACS .rtp file. Write a GROMACS .rtp file.
...@@ -82,9 +86,6 @@ class ForceField: ...@@ -82,9 +86,6 @@ class ForceField:
else: else:
return any(map(recurse, iterable)) return any(map(recurse, iterable))
def strip_polymer_bonds(bonds, char):
return [bond for bond in bonds if not any_starts_with(bond, char)]
def write_residue(name, rtp, strip=None, prepend=""): def write_residue(name, rtp, strip=None, prepend=""):
print("[ {0} ]".format(prepend + name), file=rtp) print("[ {0} ]".format(prepend + name), file=rtp)
...@@ -97,24 +98,16 @@ class ForceField: ...@@ -97,24 +98,16 @@ class ForceField:
needs_terminal_entry = [False, False] needs_terminal_entry = [False, False]
bond_tmp = bonds.get_bond_lengths(name, with_constr=True) get_bond_functions = [("bonds", functools.partial(bonds.get_bond_lengths, with_constr=True)),
if strip is not None: ("angles", bonds.get_bond_angles),
bond_tmp = strip_polymer_bonds(bond_tmp, strip) ("dihedrals", bonds.get_bond_dihedrals)]
write_bond_angle_dih(bond_tmp, "bonds", rtp)
needs_terminal_entry[0] |= any_starts_with(bond_tmp, "-")
needs_terminal_entry[1] |= any_starts_with(bond_tmp, "+")
bond_tmp = bonds.get_bond_angles(name)
if strip is not None:
bond_tmp = strip_polymer_bonds(bond_tmp, strip)
write_bond_angle_dih(bond_tmp, "angles", rtp)
needs_terminal_entry[0] |= any_starts_with(bond_tmp, "-")
needs_terminal_entry[1] |= any_starts_with(bond_tmp, "+")
bond_tmp = bonds.get_bond_dihedrals(name) for get_bond in get_bond_functions:
bond_tmp = get_bond[1](name)
if strip is not None: if strip is not None:
bond_tmp = strip_polymer_bonds(bond_tmp, strip) bond_tmp = [bond for bond in bond_tmp if not any_starts_with(bond, strip)]
write_bond_angle_dih(bond_tmp, "dihedrals", rtp, multiplicity=1) write_bond_angle_dih(bond_tmp, get_bond[0], rtp,
multiplicity=1 if get_bond[0] == "dihedrals" else None)
needs_terminal_entry[0] |= any_starts_with(bond_tmp, "-") needs_terminal_entry[0] |= any_starts_with(bond_tmp, "-")
needs_terminal_entry[1] |= any_starts_with(bond_tmp, "+") needs_terminal_entry[1] |= any_starts_with(bond_tmp, "+")
...@@ -140,11 +133,11 @@ class ForceField: ...@@ -140,11 +133,11 @@ class ForceField:
if needs_terminal_entry[1]: if needs_terminal_entry[1]:
write_residue(mol, rtp, strip="+", prepend="C") write_residue(mol, rtp, strip="+", prepend="C")
c_terms.add(mol) c_terms.add(mol)
if needs_terminal_entry[0]: if all(needs_terminal_entry):
write_residue(mol, rtp, strip=("-", "+"), prepend="2") write_residue(mol, rtp, strip=("-", "+"), prepend="2")
both_terms.add(mol) both_terms.add(mol)
self._write_r2b(filename, n_terms, c_terms, both_terms) return n_terms, c_terms, both_terms
def _write_r2b(self, filename, n_terms, c_terms, both_terms): def _write_r2b(self, filename, n_terms, c_terms, both_terms):
with open(os.path.join(self.dirname, filename + ".r2b"), "w") as r2b: with open(os.path.join(self.dirname, filename + ".r2b"), "w") as r2b:
......
...@@ -65,8 +65,7 @@ def main(args, config): ...@@ -65,8 +65,7 @@ def main(args, config):
bonds.boltzmann_invert(progress=True) bonds.boltzmann_invert(progress=True)
if config.output_forcefield: if config.output_forcefield:
logger.info("Creating GROMACS forcefield directory") logger.info("Creating GROMACS forcefield directory")
ff = ForceField(config.output_name) ForceField(config.output_name).write(config.output_name, mapping, bonds)
ff.write_rtp(config.output_name, mapping, bonds)
logger.info("GROMACS forcefield directory created") logger.info("GROMACS forcefield directory created")
else: else:
bonds.write_itp(config.output_name + ".itp", mapping=mapping) bonds.write_itp(config.output_name + ".itp", mapping=mapping)
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment