Merge branch 'workflow_parameterization' of https://github.com/shirts…

…group/terphenyl_simulations into workflow_parameterization

environment.yml

@@ -1,33 +1,39 @@
-name: terphenyl
+name: ts_env
 channels:
   - conda-forge
-  - omnia
-  - openeye
 dependencies:
-  - python
-  - pip
+  - python=3.11
+  - ca-certificates
+  - openssl
+  - ambertools
+  - openeye-toolkits
+  - psi4
+  - dftd3-python
+  - xtb-python
+  - mbuild=0.14.2
+  - tqdm
   - pytest
   - rdkit
   - pyyaml
-  - mbuild=0.14.2
   - mdtraj
   - numpy
   - matplotlib
   - scikit-learn
-  - tqdm
   - mdanalysis
   - signac
   - signac-flow
   - panedr
   - pandas
   - seaborn
-  - openff-toolkit
-  - mbuild ==0.14.2
-  - rdkit
-  - ca-certificates
-  - certifi
-  - openssl
-  - openbabel
-  - openeye-toolkits
   - gromacs
-prefix: /home/lenny/miniforge3/envs/terphenyl
+  - openbabel
+  - certifi
+  - openmm
+  - openff-units
+  - cachetools
+  - xmltodict
+  - openff-toolkit
+  - openff-bespokefit
+  - openff-interchange
+  - qcportal[version='<0.50']
+prefix: /home/lenny/miniforge3/envs/ts_env

examples/generate_bespoke_fit_parameters/mom/mom_bespoke.ipynb

@@ -62,9 +62,7 @@
    "cell_type": "code",
    "execution_count": 3,
    "id": "e018bd1b",
-   "metadata": {
-    "scrolled": false
-   },
+   "metadata": {},
    "outputs": [
     {
      "name": "stdout",
@@ -997,7 +995,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.10"
+   "version": "3.11.10"
   }
  },
  "nbformat": 4,

examples/generate_bespoke_fit_parameters/pom/pom_bespoke.ipynb

@@ -66,9 +66,7 @@
    "cell_type": "code",
    "execution_count": 3,
    "id": "e018bd1b",
-   "metadata": {
-    "scrolled": false
-   },
+   "metadata": {},
    "outputs": [
     {
      "name": "stdout",
@@ -1374,7 +1372,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.15"
+   "version": "3.11.8"
   }
  },
  "nbformat": 4,

simulations/terphenyl_mop/tetramer_remd/t_250_450/bespoke/mop_tetramer.build

@@ -6,4 +6,4 @@ residue_name : MOP
 cap_smiles : {upper : CO, lower : CC(C)(C)OC=O}
 foldamer_length : 4
 system : {solvent : TCM, solvent_smile_str : CCl, n_solvent : 634, box_size : 55.4}
-ff_method : openff-foldamer
+ff_method : bespoke-foldamer

terphenyl_simulations/__init__.py

@@ -10,10 +10,11 @@
 import terphenyl_simulations.edit_conf
 import terphenyl_simulations.build
 import terphenyl_simulations.analysis_workflows.utils
-import terphenyl_simulations.force_fields
+import terphenyl_simulations.assign_parameters
 import terphenyl_simulations.gromacs_wrapper
 import terphenyl_simulations.utils
 import terphenyl_simulations.remd_utils
+import terphenyl_simulations.topology_manager
 
 
 # Handle versioneer

terphenyl_simulations/assign_parameters.py

@@ -0,0 +1,250 @@
+from .utils import renumber_pdb_atoms, make_path
+from abc import ABC, abstractmethod
+from openmm import app
+from openff.toolkit.topology import FrozenMolecule, Molecule, Topology
+from openff.toolkit.typing.engines.smirnoff import ForceField
+from openff.interchange.components.interchange import Interchange
+from openff.bespokefit.executor import BespokeExecutor, BespokeWorkerConfig, wait_until_complete
+from openff.bespokefit.workflows import BespokeWorkflowFactory
+from openff.fragmenter.fragment import WBOFragmenter
+from openff.bespokefit.schema.targets import TorsionProfileTargetSchema
+from openff.qcsubmit.common_structures import QCSpec
+from openff.bespokefit.utilities.smirks import SMIRKSettings
+from openff.bespokefit.schema.optimizers import ForceBalanceSchema
+from openff.bespokefit.schema.smirnoff import  ProperTorsionHyperparameters
+from .topology_manager import TopologyManager
+import pdb
+import os
+import sys
+import yaml
+import subprocess
+
+
+class OFFMethod(ABC):
+    @classmethod
+    @abstractmethod
+    def assign_parameters(self):
+        pass
+
+
+class FoldamerOFFBespoke(OFFMethod):
+    def __init__(
+        self, mol_file, pdb_file, build_file = None, output_file=None, path="", ff_str="openff-2.0.0"
+    ):
+        if type(mol_file) is not str:
+            print(
+                "FoldamerOFFDefault method does cannot process multiple molecules, try"
+                + "SystemOFFDefault method instead."
+            )
+            sys.exit()
+        if not os.path.isdir(path):
+            make_path(path)
+        self.molecule = Molecule.from_file(mol_file)
+        self.name = output_file
+        self.build_file_yml = build_file
+        if output_file is None:
+            self.name = mol_file.split("/")[-1].split(".mol")[0]
+        self.path = path
+        pdb_path = pdb_file.split(".pdb")[0]
+        renumber_pdb_atoms(pdb_file, os.path.join(path, pdb_path + "_renum.pdb"))
+        self.pdb_file = app.PDBFile(os.path.join(path, pdb_path + "_renum.pdb"))
+        self.omm_topology = self.pdb_file.topology
+        self.off_topology = Topology.from_openmm(
+            self.omm_topology, unique_molecules=[self.molecule]
+        )
+        self.force_field = ff_str + ".offxml"
+        self.topology_manager = TopologyManager()
+
+    def assign_parameters(self):
+        self.generate_trimer_molecule(self.build_file_yml)
+        self.assign_trimer_partial_charges()
+        self.setup_bespoke_fit_executor()
+
+
+    def generate_trimer_molecule(self, build_file_yml):
+        from .build import FoldamerBuilder
+
+        with open(build_file_yml, "r") as f:
+            build_params = yaml.safe_load(f)
+        build_params["foldamer_length"] = 3
+        self.trimer_buildfile = build_params["structure_file"].split("_")[0] + "_trimer.build"
+        build_params["structure_file"] = build_params["structure_file"].split("_")[0] + "_trimer"
+        with open(self.trimer_buildfile, "w") as wf:
+            yaml.dump(build_params, wf)
+        foldamer_builder = FoldamerBuilder(self.trimer_buildfile)
+        foldamer_builder.get_foldamer()
+
+    def assign_trimer_partial_charges(self):
+        with open(self.trimer_buildfile, "r") as f:
+            trimer_build_params = yaml.safe_load(f)
+        self.trimer_molecule = Molecule.from_file(trimer_build_params["structure_file"] + ".mol")
+        pdb_file = trimer_build_params["structure_file"] + ".pdb"
+        renumber_pdb_atoms(pdb_file, os.path.join(self.path, pdb_file + "_renum.pdb"))
+        self.trimer_pdb  = app.PDBFile(os.path.join(self.path, pdb_file + "_renum.pdb"))
+
+        self.topology_manager.load()
+        self.trimer_sdf_file = trimer_build_params["structure_file"] + ".sdf"
+        if not self.topology_manager.check_file_type(self.trimer_buildfile, "molecule", "sdf"):
+            self.trimer_molecule.assign_partial_charges(partial_charge_method = "am1bcc")
+            self.trimer_molecule.to_file(self.trimer_sdf_file, file_format="sdf")
+            self.topology_manager.add_structure(self.trimer_sdf_file, self.trimer_buildfile, "molecule")
+        else:
+            self.topology_manager.get_structure(self.trimer_buildfile, "molecule", self.path, filetype="sdf")
+
+    def setup_bespoke_fit_executor(self, 
+                                   n_fragmenter_workers = 4,
+                                   n_qc_compute_workers = 4,
+                                   n_optimizer_workers = 4,
+                            ):
+
+        # Keep Bespoke Executor Files
+        subprocess.run(["BEFLOW_KEEP_TMP_FILES=True"], shell=True)
+
+        # Setup Executor object
+        self.bespoke_fit_executor = BespokeExecutor(
+            n_fragmenter_workers = n_fragmenter_workers,
+            n_qc_compute_workers = n_qc_compute_workers,
+            n_optimizer_workers = n_optimizer_workers,
+            launch_redis_if_unavailable = True
+        )
+
+        # Setup Workflow
+        self.factory = BespokeWorkflowFactory()
+        self.factory.target_torsion_smirks = ['[!#1]~[!$(*#*)&!D1:1]-,=;!@[!$(*#*)&!D1:2]~[!#1]']
+        self.factory.fragmentation_engine = WBOFragmenter()
+        self.factory.target_templates = [TorsionProfileTargetSchema()]
+        self.factory.default_qc_specs = [
+            QCSpec(
+                method="gfn2xtb",
+                basis=None,
+                program="xtb",
+                spec_name="xtb",
+                spec_description="gfn2xtb",
+            )
+        ]
+        self.factory.smirk_settings = SMIRKSettings(
+            expand_torsion_terms=True,
+            generate_bespoke_terms=True,
+        )
+        self.factory.initial_force_field = self.force_field
+        self.factory.optimizer = ForceBalanceSchema()
+        self.factory.parameter_hyperparameters = [ProperTorsionHyperparameters()]
+        self.factory.to_file('bespoke_flow.json')
+
+#     def run_bespoke_fit_executor(self):
+        trimer_molecule = Molecule.from_file(self.trimer_sdf_file)
+        bespoke_workflow_schema = self.factory.optimization_schema_from_molecule(trimer_molecule)
+
+        with self.bespoke_fit_executor:
+            task_id = self.bespoke_fit_executor.submit(bespoke_workflow_schema)
+            output = wait_until_complete(task_id)
+
+
+class FoldamerOFFDefault(OFFMethod):
+    def __init__(
+        self, mol_file, pdb_file, output_file=None, path="", ff_str="openff-2.0.0"
+    ):
+        if type(mol_file) is not str:
+            print(
+                "FoldamerOFFDefault method does cannot process multiple molecules, try"
+                + "SystemOFFDefault method instead."
+            )
+            sys.exit()
+        if not os.path.isdir(path):
+            make_path(path)
+        self.molecule = Molecule.from_file(mol_file)
+        self.name = output_file
+        if output_file is None:
+            self.name = mol_file.split("/")[-1].split(".mol")[0]
+        self.path = path
+        pdb_path = pdb_file.split(".pdb")[0]
+        renumber_pdb_atoms(pdb_file, os.path.join(path, pdb_path + "_renum.pdb"))
+        self.pdb_file = app.PDBFile(os.path.join(path, pdb_path + "_renum.pdb"))
+        self.omm_topology = self.pdb_file.topology
+        self.off_topology = Topology.from_openmm(
+            self.omm_topology, unique_molecules=[self.molecule]
+        )
+        self.force_field = ForceField(ff_str + ".offxml")
+
+    def assign_parameters(self, charge_method="am1bcc"):
+        self._get_partial_charges(charge_method)
+        top_file, gro_file = self._generate_ff_topologies()
+        return top_file, gro_file, self.sdf_file
+
+    def _get_partial_charges(self, method="am1bcc"):
+        self.sdf_file = os.path.join(self.path, self.name + "_charges.sdf")
+        if not os.path.exists(self.sdf_file):
+            # Expensive step
+            self.molecule.assign_partial_charges(partial_charge_method=method)
+            self.molecule.to_file(self.sdf_file, file_format="sdf")
+        else:
+            self.molecule = Molecule.from_file(self.sdf_file)
+
+    def _generate_ff_topologies(self):
+        interchange = Interchange.from_smirnoff(
+            force_field=self.force_field,
+            topology=self.off_topology,
+            charge_from_molecules=[self.molecule],
+        )
+        interchange.positions = self.pdb_file.getPositions()
+
+        top_file = os.path.join(self.path, self.name + "_openff-2.0.0.top")
+        gro_file = os.path.join(self.path, self.name + "_openff-2.0.0.gro")
+        interchange.to_top(top_file)
+        interchange.to_gro(gro_file)
+
+        return top_file, gro_file
+
+
+class SystemOFFDefault(OFFMethod):
+    def __init__(
+        self,
+        system_molecules_list,
+        charge_files,
+        system_pdb,
+        path="",
+        ff_str="openff-2.0.0",
+    ):
+        if not os.path.isdir(path):
+            make_path(path)
+        self.path = path
+        self.molecules = []
+        self.charges = []
+        for molecule, charge_file in zip(system_molecules_list, charge_files):
+            off_molecule = Molecule.from_file(molecule)
+            # If we have partial charges in a file use it
+            if charge_file != None and os.path.exists(charge_file):
+                print("Getting charges for", molecule, "from", charge_file)
+                molecule_charges = Molecule.from_file(charge_file).partial_charges
+                off_molecule.partial_charges = molecule_charges
+                off_molecule.perceive_residues()
+            self.molecules.append(off_molecule)
+            self.charges.append(off_molecule.partial_charges != None)
+        self.pdb_file = app.PDBFile(system_pdb)
+        print([mol.name for mol in self.molecules])
+        print(self.pdb_file.topology)
+        self.off_topology = Topology.from_openmm(
+            self.pdb_file.topology, unique_molecules=self.molecules
+        )
+        self.force_field = ForceField(ff_str + ".offxml")
+
+    def _generate_ff_topology(self):
+        interchange = Interchange.from_smirnoff(
+            force_fild=self.force_field,
+            topology=self.off_topology,
+            charge_from_molecules=[
+                mol for mol, charge in zip(self.molecules, self.charges) if charge
+            ],
+        )
+        interchange.positions = self.pdb_file.getPositions()
+
+        top_file = os.path.join(self.path, self.name + "_openff-2.0.0.top")
+        gro_file = os.path.join(self.path, self.name + "_openff-2.0.0.gro")
+        interchange.to_top(top_file)
+        interchange.to_gro(gro_file)
+
+        return top_file, gro_file
+
+    def assign_parameters(self):
+        top_file, gro_file = self._generate_ff_topologies()
+        return top_file, gro_file