From b73a532294d0f60163f3b51ca17c04f56e749f3a Mon Sep 17 00:00:00 2001 From: andrewtarzia Date: Tue, 21 Nov 2023 21:46:55 +0100 Subject: [PATCH] Refactor and add Martini system. --- src/cgexplore/assigned_system.py | 383 ++++++++++++++++++++++--------- 1 file changed, 273 insertions(+), 110 deletions(-) diff --git a/src/cgexplore/assigned_system.py b/src/cgexplore/assigned_system.py index 8263008c..363cfff6 100644 --- a/src/cgexplore/assigned_system.py +++ b/src/cgexplore/assigned_system.py @@ -7,7 +7,6 @@ """ -import os import pathlib from dataclasses import dataclass @@ -16,108 +15,14 @@ from .beads import CgBead, get_cgbead_from_element from .errors import ForcefieldUnavailableError, ForcefieldUnitError +from .martini import MartiniTopology, get_martini_mass_by_type from .utilities import ( cosine_periodic_angle_force, custom_excluded_volume_force, ) -@dataclass(frozen=True) -class AssignedSystem: - molecule: stk.Molecule - force_field_terms: dict[str, tuple] - system_xml: pathlib.Path - topology_xml: pathlib.Path - bead_set: dict[str, CgBead] - vdw_bond_cutoff: int - mass: float = 10 - - def _get_topology_xml_string(self, molecule: stk.Molecule) -> str: - ff_str = "\n\n" - - at_str = " \n" - re_str = " \n" - re_str += ' \n' - - present_beads = {} - for atom in molecule.get_atoms(): - aestring = atom.__class__.__name__ - aid = atom.get_id() - acgbead = get_cgbead_from_element( - estring=aestring, - bead_set=self.bead_set, - ) - atype = acgbead.bead_type - aclass = acgbead.bead_class - - if atype not in present_beads: - present_beads[atype] = acgbead - at_str += ( - f' \n' - ) - - re_str += f' \n' - - for bond in molecule.get_bonds(): - a1id = bond.get_atom1().get_id() - a2id = bond.get_atom2().get_id() - - re_str += f' \n' - - at_str += " \n\n" - re_str += " \n" - re_str += " \n\n" - - ff_str += at_str - ff_str += re_str - - ff_str += "\n" - return ff_str - - def _write_topology_xml(self, molecule: stk.Molecule) -> None: - ff_str = self._get_topology_xml_string(molecule) - - with open(self.topology_xml, "w") as f: - f.write(ff_str) - - def get_openmm_topology(self) -> app.topology.Topology: - topology = app.topology.Topology() - chain = topology.addChain() - residue = topology.addResidue(name="ALL", chain=chain) - - atoms_added = {} - for atom in self.molecule.get_atoms(): - a_id = atom.get_id() - a_estring = atom.__class__.__name__ - a_element = app.element.Element.getByAtomicNumber( - atom.get_atomic_number() - ) - a_cgbead = get_cgbead_from_element( - estring=a_estring, - bead_set=self.bead_set, - ) - - omm_atom = topology.addAtom( - name=a_cgbead.bead_type, - element=a_element, - residue=residue, - id=str(a_id), - ) - atoms_added[a_id] = omm_atom - - for bond in self.molecule.get_bonds(): - a1_id = bond.get_atom1().get_id() - a2_id = bond.get_atom2().get_id() - - topology.addBond( - atom1=atoms_added[a1_id], - atom2=atoms_added[a2_id], - ) - - return topology - +class ForcedSystem: def _available_forces(self, force_type: str) -> openmm.Force: available = { "HarmonicBondForce": openmm.HarmonicBondForce(), @@ -135,6 +40,8 @@ def _add_bonds(self, system: openmm.System) -> openmm.System: forces = self.force_field_terms["bond"] force_types = {i.force for i in forces} for force_type in force_types: + if "Martini" in force_type: + continue force_function = self._available_forces(force_type) system.addForce(force_function) for assigned_force in forces: @@ -163,6 +70,8 @@ def _add_angles(self, system: openmm.System) -> openmm.System: forces = self.force_field_terms["angle"] force_types = {i.force for i in forces} for force_type in force_types: + if "Martini" in force_type: + continue force_function = self._available_forces(force_type) system.addForce(force_function) for assigned_force in forces: @@ -170,15 +79,6 @@ def _add_angles(self, system: openmm.System) -> openmm.System: continue try: if force_type == "CosinePeriodicAngleForce": - print( - [ - assigned_force.angle_k.value_in_unit( - openmm.unit.kilojoule / openmm.unit.mole - ), - assigned_force.n, - assigned_force.b, - ], - ) force_function.addAngle( assigned_force.atom_ids[0], assigned_force.atom_ids[1], @@ -217,6 +117,8 @@ def _add_torsions(self, system: openmm.System) -> openmm.System: forces = self.force_field_terms["torsion"] force_types = {i.force for i in forces} for force_type in force_types: + if "Martini" in force_type: + continue force_function = self._available_forces(force_type) system.addForce(force_function) for assigned_force in forces: @@ -289,11 +191,104 @@ def _add_forces(self, system: openmm.System) -> openmm.System: system = self._add_torsions(system) return self._add_nonbondeds(system) - def get_openmm_system(self) -> openmm.System: - if os.path.exists(self.system_xml): - with open(self.system_xml) as f: - return openmm.XmlSerializer.deserialize(f.read()) +@dataclass(frozen=True) +class AssignedSystem(ForcedSystem): + molecule: stk.Molecule + force_field_terms: dict[str, tuple] + system_xml: pathlib.Path + topology_xml: pathlib.Path + bead_set: dict[str, CgBead] + vdw_bond_cutoff: int + mass: float = 10 + + def _get_topology_xml_string(self, molecule: stk.Molecule) -> str: + ff_str = "\n\n" + + at_str = " \n" + re_str = " \n" + re_str += ' \n' + + present_beads = {} + for atom in molecule.get_atoms(): + aestring = atom.__class__.__name__ + aid = atom.get_id() + acgbead = get_cgbead_from_element( + estring=aestring, + bead_set=self.bead_set, + ) + atype = acgbead.bead_type + aclass = acgbead.bead_class + + if atype not in present_beads: + present_beads[atype] = acgbead + at_str += ( + f' \n' + ) + + re_str += f' \n' + + for bond in molecule.get_bonds(): + a1id = bond.get_atom1().get_id() + a2id = bond.get_atom2().get_id() + + re_str += f' \n' + + at_str += " \n\n" + re_str += " \n" + re_str += " \n\n" + + ff_str += at_str + ff_str += re_str + + ff_str += "\n" + return ff_str + + def _write_topology_xml(self, molecule: stk.Molecule) -> None: + ff_str = self._get_topology_xml_string(molecule) + + with open(self.topology_xml, "w") as f: + f.write(ff_str) + + def get_openmm_topology(self) -> app.topology.Topology: + topology = app.topology.Topology() + chain = topology.addChain() + residue = topology.addResidue(name="ALL", chain=chain) + + atoms_added = {} + for atom in self.molecule.get_atoms(): + a_id = atom.get_id() + a_estring = atom.__class__.__name__ + a_element = app.element.Element.getByAtomicNumber( + atom.get_atomic_number() + ) + a_cgbead = get_cgbead_from_element( + estring=a_estring, + bead_set=self.bead_set, + ) + + omm_atom = topology.addAtom( + name=a_cgbead.bead_type, + element=a_element, + residue=residue, + id=str(a_id), + ) + atoms_added[a_id] = omm_atom + + for bond in self.molecule.get_bonds(): + a1_id = bond.get_atom1().get_id() + a2_id = bond.get_atom2().get_id() + + topology.addBond( + atom1=atoms_added[a1_id], + atom2=atoms_added[a2_id], + ) + + return topology + + def get_openmm_system(self) -> openmm.System: self._write_topology_xml(self.molecule) forcefield = app.ForceField(self.topology_xml) topology = self.get_openmm_topology() @@ -304,3 +299,171 @@ def get_openmm_system(self) -> openmm.System: f.write(openmm.XmlSerializer.serialize(system)) return system + + +@dataclass(frozen=True) +class MartiniSystem(ForcedSystem): + """ + Assign a system using martini_openmm. + + """ + + molecule: stk.Molecule + force_field_terms: dict[str, tuple] + system_xml: pathlib.Path + topology_itp: pathlib.Path + vdw_bond_cutoff: int + bead_set: dict[str, CgBead] + + def _get_atoms_string( + self, + molecule: stk.Molecule, + molecule_name: str, + ) -> str: + atoms_string = ( + "[atoms]\n" + ";nr type resnr resid atom cgnr charge mass total_charge\n" + ) + for atom in molecule.get_atoms(): + a_estring = atom.__class__.__name__ + + a_cgbead = get_cgbead_from_element( + estring=a_estring, + bead_set=self.bead_set, + ) + nr = atom.get_id() + 1 + type_ = a_cgbead.bead_type + resnr = 1 + resid = molecule_name[:4].upper() + charge = 0 + total_charge = 0 + # Charge group, set as different for all for now, like in PYRI. + cgnr = atom.get_id() + 1 + mass = get_martini_mass_by_type(type_) + atoms_string += ( + f"{nr} {type_} {resnr} {resid} {a_estring} {cgnr} {charge} " + f"{mass} {total_charge}\n" + ) + atoms_string += "\n" + return atoms_string + + def _get_bonds_string(self, molecule: stk.Molecule) -> str: + string = "[bonds]\n; i j funct length\n" + forces = self.force_field_terms["bond"] + for assigned_force in forces: + force_type = assigned_force.force + if force_type != "MartiniDefinedBond": + continue + length = assigned_force.bond_r.value_in_unit(openmm.unit.nanometer) + k = assigned_force.bond_k.value_in_unit( + openmm.unit.kilojoule + / openmm.unit.mole + / openmm.unit.nanometer**2 + ) + string += ( + f" {assigned_force.atom_ids[0]+1} " + f"{assigned_force.atom_ids[1]+1} " + f"{assigned_force.funct} " + f"{length} " + f"{k}\n" + ) + string += "\n" + return string + + def _get_angles_string(self, molecule: stk.Molecule) -> str: + string = "[angles]\n; i j k funct ref.angle force_k\n" + forces = self.force_field_terms["angle"] + + for assigned_force in forces: + force_type = assigned_force.force + if force_type != "MartiniDefinedAngle": + continue + + angle = assigned_force.angle.value_in_unit(openmm.unit.degrees) + k = assigned_force.angle_k.value_in_unit( + openmm.unit.kilojoule + / openmm.unit.mole + / openmm.unit.radian**2 + ) + string += ( + f" {assigned_force.atom_ids[0]+1} " + f"{assigned_force.atom_ids[1]+1} " + f"{assigned_force.atom_ids[2]+1} " + f"{assigned_force.funct} " + f"{angle} " + f"{k}\n" + ) + + string += "\n" + return string + + def _get_torsions_string(self, molecule: stk.Molecule) -> str: + string = "[dihedrals]\n; i j k l funct ref.angle force_k\n" + forces = self.force_field_terms["torsion"] + force_types = {i.force for i in forces} + + for force_type in force_types: + if force_type != "MartiniDefinedTorsion": + continue + print(force_type) + raise SystemExit() + string += "\n" + return string + + def _get_contraints_string(self, molecule: stk.Molecule) -> str: + string = "[constraints]\n; i j funct length\n" + for constraint in self.force_field_terms["constraints"]: + string += ( + f" {constraint[0]} {constraint[1]} {constraint[2]} " + f"{constraint[3]} {constraint[4]}" + ) + string += "\n" + return string + + def _get_exclusions_string(self, molecule: stk.Molecule) -> str: + string = "[exclusions]\n; i j funct length\n" + for constraint in self.force_field_terms["constraints"]: + string += ( + f" {constraint[0]} {constraint[1]} {constraint[2]} " + f"{constraint[3]} {constraint[4]}" + ) + string += "\n" + return string + + def _write_topology_itp(self, molecule: stk.Molecule) -> None: + molecule_name = self.topology_itp.name.replace(".itp", "") + string = ( + f";;; {molecule_name}\n" + "[moleculetype]\n" + "; Name nrexcl\n" + f"{molecule_name} {self.vdw_bond_cutoff}\n\n" + ) + + atoms_string = self._get_atoms_string(molecule, molecule_name) + bonds_string = self._get_bonds_string(molecule) + angles_string = self._get_angles_string(molecule) + torsions_string = self._get_torsions_string(molecule) + constraints_string = self._get_contraints_string(molecule) + + string += atoms_string + string += bonds_string + string += angles_string + string += torsions_string + string += constraints_string + + with open(self.topology_itp, "w") as f: + f.write(string) + + def get_openmm_topology(self) -> app.topology.Topology: + self._write_topology_itp(self.molecule) + return MartiniTopology(self.topology_itp).get_openmm_topology() + + def get_openmm_system(self) -> openmm.System: + self._write_topology_itp(self.molecule) + topology = MartiniTopology(self.topology_itp) + system = topology.get_openmm_system() + + system = self._add_forces(system) + with open(self.system_xml, "w") as f: + f.write(openmm.XmlSerializer.serialize(system)) + return system