cleaning up a little bit

benchmarks/subsampled/gen_systems.py

@@ -1,18 +1,19 @@
-import json
 import gzip
+import json
 import pathlib
-from rdkit import Chem
+
+from gufe import ChemicalSystem, SmallMoleculeComponent
 from gufe.tokenization import JSON_HANDLER
-from gufe import SmallMoleculeComponent, ChemicalSystem
 from openff.toolkit import Molecule
 from pontibus.components import ExtendedSolventComponent
+from rdkit import Chem
 
 
 def add_chemical_systems(
     sdffile: str,
     dataset_name: str,
     solvents: dict[str, SmallMoleculeComponent],
-    systems: list[ChemicalSystem]
+    systems: list[ChemicalSystem],
 ) -> None:
     """
     Add Solute + Solvent ChemicalSystems to running list.
@@ -31,21 +32,25 @@ def add_chemical_systems(
     """
     for i, rdmol in enumerate(Chem.SDMolSupplier(sdffile, removeHs=False)):
         offmol = Molecule.from_rdkit(rdmol)
-        offmol.assign_partial_charges(partial_charge_method='am1bccelf10')
-        solvent_smi = rdmol.GetProp('solvent')
+        offmol.assign_partial_charges(partial_charge_method="am1bccelf10")
+        solvent_smi = rdmol.GetProp("solvent")
         if solvent_smi not in solvents.keys():
             solvent_offmol = Molecule.from_smiles(solvent_smi)
             solvent_offmol.generate_conformers()
-            solvent_offmol.assign_partial_charges(partial_charge_method='am1bccelf10')
+            solvent_offmol.assign_partial_charges(partial_charge_method="am1bccelf10")
             solvents[solvent_smi] = SmallMoleculeComponent.from_openff(solvent_offmol)
 
-        systems.append(ChemicalSystem(
-            {
-                'solute': SmallMoleculeComponent.from_openff(offmol),
-                'solvent': ExtendedSolventComponent(solvent_molecule=solvents[solvent_smi]),
-            },
-            name=f"molecule{i}_{dataset_name}"
-        ))
+        systems.append(
+            ChemicalSystem(
+                {
+                    "solute": SmallMoleculeComponent.from_openff(offmol),
+                    "solvent": ExtendedSolventComponent(
+                        solvent_molecule=solvents[solvent_smi]
+                    ),
+                },
+                name=f"molecule{i}_{dataset_name}",
+            )
+        )
 
 
 def store_chemical_systems(systems: list[ChemicalSystem], outdir: pathlib.Path):
@@ -58,17 +63,17 @@ def store_chemical_systems(systems: list[ChemicalSystem], outdir: pathlib.Path):
       List of ChemicalSystems to store to file.
     """
     for system in systems:
-        with gzip.open(outdir / f"{system.name}_chemicalsystem.gz", 'wt') as zipfile:
+        with gzip.open(outdir / f"{system.name}_chemicalsystem.gz", "wt") as zipfile:
             json.dump(system.to_dict(), zipfile, cls=JSON_HANDLER.encoder)
 
 
 if __name__ == "__main__":
     solvents: dict[str, SmallMoleculeComponent] = {}
     systems: list[ChemicalSystem] = []
 
-    add_chemical_systems('sub_sampled_fsolv.sdf', 'fsolv', solvents, systems)
-    add_chemical_systems('sub_sampled_mnsol.sdf', 'mnsol', solvents, systems)
+    add_chemical_systems("sub_sampled_fsolv.sdf", "fsolv", solvents, systems)
+    add_chemical_systems("sub_sampled_mnsol.sdf", "mnsol", solvents, systems)
 
-    outdir = pathlib.Path('chemicalsystems')
+    outdir = pathlib.Path("chemicalsystems")
     outdir.mkdir(exist_ok=True)
-    store_chemical_systems(systems, outdir)
+    store_chemical_systems(systems, outdir)

benchmarks/subsampled/gen_transforms_2.0.0_single_repeat.py

@@ -1,26 +1,21 @@
-import pathlib
-import json
 import gzip
-from openff.units import unit
+import json
+import pathlib
+
+from gufe import AlchemicalNetwork, ChemicalSystem, Transformation
 from gufe.tokenization import JSON_HANDLER
-from gufe import (
-    AlchemicalNetwork,
-    ChemicalSystem,
-    Transformation
-)
+from openff.toolkit import Molecule
+from openff.units import unit
 from pontibus.protocols.solvation import ASFEProtocol
 
 
 def deserialize_system(file: pathlib.Path):
-    with gzip.open(file, 'r') as f:
-        d = json.loads(
-            f.read().decode(),
-            cls=JSON_HANDLER.decoder
-        )
+    with gzip.open(file, "r") as f:
+        d = json.loads(f.read().decode(), cls=JSON_HANDLER.decoder)
     return ChemicalSystem.from_dict(d)
 
 
-def get_settings():
+def get_water_settings():
     settings = ASFEProtocol.default_settings()
     settings.protocol_repeats = 1
     settings.solvent_forcefield_settings.forcefields = [
@@ -32,27 +27,52 @@ def get_settings():
     ]
     settings.solvent_simulation_settings.time_per_iteration = 5 * unit.picosecond
     settings.vacuum_simulation_settings.time_per_iteration = 5 * unit.picosecond
+    settings.vacuum_engine_settings.compute_platform = "CPU"
+    settings.solvent_engine_settings.compute_platform = "CUDA"
+    settings.solvation_settings.box_shape = "dodecahedron"
+    return settings
+
+
+def get_nonwater_settings():
+    settings = ASFEProtocol.default_settings()
+    settings.protocol_repeats = 1
+    settings.solvent_forcefield_settings.forcefields = [
+        "openff-2.0.0.offxml",
+        "tip3p.offxml",
+    ]
+    settings.vacuum_forcefield_settings.forcefields = [
+        "openff-2.0.0.offxml",
+    ]
+    settings.solvent_simulation_settings.time_per_iteration = 5 * unit.picosecond
+    settings.vacuum_simulation_settings.time_per_iteration = 5 * unit.picosecond
+    settings.vacuum_engine_settings.compute_platform = "CUDA"
+    settings.solvent_engine_settings.compute_platform = "CUDA"
+    settings.solvation_settings.box_shape = "dodecahedron"
+    settings.solvation_settings.assign_solvent_charges = True
     return settings
 
 
 def get_transformation(stateA):
-    settings = get_settings()
-    stateB = stateA.components['solvent']
+    solvent = stateA.components["solvent"]
+
+    water = Molecule.from_smiles("O")
+    if water.is_isomorphic_with(solvent.solvent_molecule.to_openff()):
+        settings = get_water_settings()
+    else:
+        settings = get_nonwater_settings()
+
+    stateB = ChemicalSystem({"solvent": solvent})
     protocol = ASFEProtocol(settings=settings)
     return Transformation(
-        stateA=stateA,
-        stateB=stateB,
-        mapping=None,
-        protocol=protocol,
-        name=stateA.name
+        stateA=stateA, stateB=stateB, mapping=None, protocol=protocol, name=stateA.name
     )
 
 
 def run(outdir: pathlib.Path):
     systems: list[ChemicalSystem] = []
 
-    systems_dir = pathlib.Path('chemicalsystems')
-    system_files = systems_dir.glob('*.gz')
+    systems_dir = pathlib.Path("chemicalsystems")
+    system_files = systems_dir.glob("*.gz")
 
     for file in system_files:
         systems.append(deserialize_system(file))
@@ -77,6 +97,6 @@ def run(outdir: pathlib.Path):
 
 
 if __name__ == "__main__":
-    outdir = pathlib.Path('2.0.0_single_repeat_inputs')
+    outdir = pathlib.Path("2.0.0_single_repeat_inputs")
     outdir.mkdir(exist_ok=False)
     run(outdir)

devtools/gen_serialized_results.py

@@ -4,34 +4,32 @@
 Generates
 - ASFEProtocol_json_results.gz
 """
+
 import gzip
 import json
 import logging
 import pathlib
 import tempfile
-from openff.toolkit import (
-    Molecule, RDKitToolkitWrapper, AmberToolsToolkitWrapper
-)
+
+import gufe
+import openfe
+from gufe.tokenization import JSON_HANDLER
+from kartograf import KartografAtomMapper
+from kartograf.atom_aligner import align_mol_shape
+from openff.toolkit import AmberToolsToolkitWrapper, Molecule, RDKitToolkitWrapper
 from openff.toolkit.utils.toolkit_registry import (
-    toolkit_registry_manager, ToolkitRegistry
+    ToolkitRegistry,
+    toolkit_registry_manager,
 )
 from openff.units import unit
-from kartograf.atom_aligner import align_mol_shape
-from kartograf import KartografAtomMapper
-import gufe
-from gufe.tokenization import JSON_HANDLER
-import openfe
-from pontibus.protocols.solvation import ASFEProtocol
 from pontibus.components import ExtendedSolventComponent
-
+from pontibus.protocols.solvation import ASFEProtocol
 
 logger = logging.getLogger(__name__)
 
 LIGA = "[H]C([H])([H])C([H])([H])C(=O)C([H])([H])C([H])([H])[H]"
 
-amber_rdkit = ToolkitRegistry(
-    [RDKitToolkitWrapper(), AmberToolsToolkitWrapper()]
-)
+amber_rdkit = ToolkitRegistry([RDKitToolkitWrapper(), AmberToolsToolkitWrapper()])
 
 
 def get_molecule(smi, name):
@@ -51,7 +49,7 @@ def execute_and_serialize(dag, protocol, simname):
             shared_basedir=workdir,
             scratch_basedir=workdir,
             keep_shared=False,
-            n_retries=3
+            n_retries=3,
         )
     protres = protocol.gather([dagres])
 
@@ -60,41 +58,42 @@ def execute_and_serialize(dag, protocol, simname):
         "uncertainty": protres.get_uncertainty(),
         "protocol_result": protres.to_dict(),
         "unit_results": {
-            unit.key: unit.to_keyed_dict()
-            for unit in dagres.protocol_unit_results
-        }
+            unit.key: unit.to_keyed_dict() for unit in dagres.protocol_unit_results
+        },
     }
 
-    with gzip.open(f"{simname}_json_results.gz", 'wt') as zipfile:
+    with gzip.open(f"{simname}_json_results.gz", "wt") as zipfile:
         json.dump(outdict, zipfile, cls=JSON_HANDLER.encoder)
 
 
 def generate_ahfe_settings():
     settings = ASFEProtocol.default_settings()
-    settings.solvent_equil_simulation_settings.equilibration_length_nvt = 10 * unit.picosecond
-    settings.solvent_equil_simulation_settings.equilibration_length = 10 * unit.picosecond
+    settings.solvent_equil_simulation_settings.equilibration_length_nvt = (
+        10 * unit.picosecond
+    )
+    settings.solvent_equil_simulation_settings.equilibration_length = (
+        10 * unit.picosecond
+    )
     settings.solvent_equil_simulation_settings.production_length = 10 * unit.picosecond
     settings.solvent_simulation_settings.equilibration_length = 10 * unit.picosecond
     settings.solvent_simulation_settings.production_length = 500 * unit.picosecond
-    settings.vacuum_equil_simulation_settings.equilibration_length = 10 * unit.picosecond
+    settings.vacuum_equil_simulation_settings.equilibration_length = (
+        10 * unit.picosecond
+    )
     settings.vacuum_equil_simulation_settings.production_length = 10 * unit.picosecond
     settings.vacuum_simulation_settings.equilibration_length = 10 * unit.picosecond
     settings.vacuum_simulation_settings.production_length = 500 * unit.picosecond
     settings.protocol_repeats = 3
-    settings.vacuum_engine_settings.compute_platform = 'CPU'
-    settings.solvent_engine_settings.compute_platform = 'CUDA'
+    settings.vacuum_engine_settings.compute_platform = "CPU"
+    settings.solvent_engine_settings.compute_platform = "CUDA"
 
     return settings
 
-    
+
 def generate_asfe_json(smc):
     protocol = ASFEProtocol(settings=generate_ahfe_settings())
-    sysA = openfe.ChemicalSystem(
-        {"ligand": smc, "solvent": ExtendedSolventComponent()}
-    )
-    sysB = openfe.ChemicalSystem(
-        {"solvent": ExtendedSolventComponent()}
-    )
+    sysA = openfe.ChemicalSystem({"ligand": smc, "solvent": ExtendedSolventComponent()})
+    sysB = openfe.ChemicalSystem({"solvent": ExtendedSolventComponent()})
 
     dag = protocol.create(stateA=sysA, stateB=sysB, mapping=None)