diff --git a/devtools/gen_serialized_results.py b/devtools/gen_serialized_results.py index c47e117..8138c0d 100644 --- a/devtools/gen_serialized_results.py +++ b/devtools/gen_serialized_results.py @@ -14,6 +14,8 @@ 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 ( ToolkitRegistry, @@ -22,11 +24,13 @@ from openff.units import unit from pontibus.components import ExtendedSolventComponent +from pontibus.protocols.relative import HybridTopProtocol 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]" +LIGB = "[H]C([H])([H])C(=O)C([H])([H])C([H])([H])C([H])([H])[H]" amber_rdkit = ToolkitRegistry([RDKitToolkitWrapper(), AmberToolsToolkitWrapper()]) @@ -111,7 +115,62 @@ def generate_asfe_json_octanol(smc): execute_and_serialize(dag, protocol, "ASFEProtocol_octanol") +def generate_rfe_settings(): + settings = HybridTopProtocol.default_settings() + settings.simulation_settings.equilibration_length = 10 * unit.picosecond + settings.simulation_settings.production_length = 250 * unit.picosecond + return settings + + +def generate_rfe_inputs(molA, molB): + a_molB = align_mol_shape(molB, ref_mol=molA) + mapper = KartografAtomMapper(atom_map_hydrogens=True) + mapping = next(mapper.suggest_mappings(molA, a_molB)) + + return mapping, molA, a_molB + + +def generate_hybridtop_rfe_vacuum_json(molA, molB): + settings = generate_rfe_settings() + settings.forcefield_settings.nonbonded_method = "nocutoff" + protocol = HybridTopProtocol(settings=settings) + + mapping, molA, molB = generate_rfe_inputs(molA, molB) + + systemA = openfe.ChemicalSystem({"ligand": molA}) + systemB = openfe.ChemicalSystem({"ligand": molB}) + + dag = protocol.create( + stateA=systemA, + stateB=systemB, + mapping=mapping, + ) + + execute_and_serialize(dag, protocol, "HybridTopProtocol_vacuum") + + +def generate_hybridtop_rfe_solvent_json(molA, molB): + settings = generate_rfe_settings() + protocol = HybridTopProtocol(settings=settings) + + mapping, molA, molB = generate_rfe_inputs(molA, molB) + + systemA = openfe.ChemicalSystem({"ligand": molA, "solvent": openfe.SolventComponent()}) + systemB = openfe.ChemicalSystem({"ligand": molB, "solvent": openfe.SolventComponent()}) + + dag = protocol.create( + stateA=systemA, + stateB=systemB, + mapping=mapping, + ) + + execute_and_serialize(dag, protocol, "HybridTopProtocol_solvent") + + if __name__ == "__main__": molA = get_molecule(LIGA, "ligandA") - # generate_asfe_json_water(molA) + molB = get_molecule(LIGB, "ligandB") + generate_asfe_json_water(molA) generate_asfe_json_octanol(molA) + generate_hybridtop_rfe_vacuum_json(molA, molB) + generate_hybridtop_rfe_solvent_json(molA, molB) diff --git a/environment.yml b/environment.yml index b7022a0..5debf74 100644 --- a/environment.yml +++ b/environment.yml @@ -3,7 +3,7 @@ channels: - conda-forge dependencies: - pip - - gufe >=1.5 + - gufe >=1.6 - openff-toolkit >0.16.0 - openff-interchange >=0.4 - openff-nagl-base >=0.3.3 @@ -20,12 +20,13 @@ dependencies: - pydantic >=2.0 - pyyaml - coverage - - cinnabar ~=0.4.0 + - cinnabar ~=0.5.0 - click - typing-extensions - openmm >=8.0.0,!=8.1.0,<8.4.0 - openmmtools >=0.25.0 - openmmforcefields + - openfe-analysis >=0.3.1 - plugcli - tqdm # Testing deps @@ -36,4 +37,4 @@ dependencies: - pip: - git+https://github.com/OpenFreeEnergy/openfe@main - git+https://github.com/OpenFreeEnergy/ofe-sphinx-theme@a45f3edd5bc3e973c1a01b577c71efa1b62a65d6 - - git+https://github.com/openforcefield/openff-interchange@sltcap + - git+https://github.com/openforcefield/openff-interchange@main diff --git a/src/pontibus/protocols/relative/__init__.py b/src/pontibus/protocols/relative/__init__.py new file mode 100644 index 0000000..4817b4f --- /dev/null +++ b/src/pontibus/protocols/relative/__init__.py @@ -0,0 +1,19 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe +""" +Run relative free energy calculations using OpenMM and OpenMMTools. +""" + +from .hybridtop_protocol import ( + HybridTopProtocol, + HybridTopProtocolResult, +) +from .hybridtop_units import HybridTopProtocolUnit +from .settings import HybridTopProtocolSettings + +__all__ = [ + "HybridTopProtocol", + "HybridTopProtocolSettings", + "HybridTopProtocolResult", + "HybridTopProtocolUnit", +] diff --git a/src/pontibus/protocols/relative/hybridtop_protocol.py b/src/pontibus/protocols/relative/hybridtop_protocol.py new file mode 100644 index 0000000..acf1d0d --- /dev/null +++ b/src/pontibus/protocols/relative/hybridtop_protocol.py @@ -0,0 +1,145 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe +"""Equilibrium Relative Free Energy methods using OpenMM and OpenMMTools in a +Perses-like manner. +This module implements the necessary methodology toolking to run calculate a +ligand relative free energy transformation using OpenMM tools and one of the +following methods: + - Hamiltonian Replica Exchange + - Self-adjusted mixture sampling + - Independent window sampling + +Acknowledgements +---------------- +This Protocol is a subclass of the OpenFE RelativeHybridTopologyProtocol. +This Protocol is based on, and leverages components originating from +the Perses toolkit (https://github.com/choderalab/perses). +""" + +import uuid + +from gufe import ChemicalSystem, ComponentMapping, ProtocolDAGResult, ProtocolUnit +from gufe.settings import ThermoSettings +from openfe.protocols.openmm_rfe.equil_rfe_methods import ( + RelativeHybridTopologyProtocol, + RelativeHybridTopologyProtocolResult, + _validate_alchemical_components, +) +from openfe.protocols.openmm_rfe.equil_rfe_settings import ( + AlchemicalSettings, + LambdaSettings, +) +from openfe.protocols.openmm_utils import system_validation +from openfe.protocols.openmm_utils.omm_settings import ( + IntegratorSettings, + MultiStateOutputSettings, + MultiStateSimulationSettings, + OpenFFPartialChargeSettings, + OpenMMEngineSettings, +) +from openff.units import unit + +from pontibus.protocols.relative.hybridtop_units import HybridTopProtocolUnit +from pontibus.protocols.relative.settings import HybridTopProtocolSettings +from pontibus.utils.settings import InterchangeFFSettings, PackmolSolvationSettings + + +class HybridTopProtocolResult(RelativeHybridTopologyProtocolResult): + """ + Results class for the HybridTopologyProtocol class. + Inherits from + :class:`openfe.protocols.openmm_rfe.RelativeHybridTopologyProtocolResult`. + """ + + +class HybridTopProtocol(RelativeHybridTopologyProtocol): + """ + Relative Free Energy calculations using OpenMM and OpenMMTools. + + Based on `Perses `_ + + See Also + -------- + :mod:`openfe.protocols` + :class:`openfe.protocols.openmm_rfe.RelativeHybridTopologyProtocol` + :class:`pontibus.protocols.relative.HybridTopSettings` + :class:`pontibus.protocols.relative.HybridTopResult` + :class:`pontibus.protocols.relative.HybridTopProtocolUnit` + """ + + result_cls = HybridTopProtocolResult + _settings_cls = HybridTopProtocolSettings + _settings: HybridTopProtocolSettings + + @classmethod + def _default_settings(cls): + """A dictionary of initial settings for this creating this Protocol + These settings are intended as a suitable starting point for creating + an instance of this protocol. It is recommended, however that care is + taken to inspect and customize these before performing a Protocol. + Returns + ------- + Settings + a set of default settings + """ + return HybridTopProtocolSettings( + protocol_repeats=3, + forcefield_settings=InterchangeFFSettings(), + thermo_settings=ThermoSettings( + temperature=298.15 * unit.kelvin, + pressure=1 * unit.bar, + ), + partial_charge_settings=OpenFFPartialChargeSettings(), + solvation_settings=PackmolSolvationSettings(), + alchemical_settings=AlchemicalSettings(softcore_LJ="gapsys"), + lambda_settings=LambdaSettings(), + simulation_settings=MultiStateSimulationSettings( + equilibration_length=1.0 * unit.nanosecond, + production_length=5.0 * unit.nanosecond, + ), + engine_settings=OpenMMEngineSettings(), + integrator_settings=IntegratorSettings(), + output_settings=MultiStateOutputSettings(), + ) + + def _create( + self, + stateA: ChemicalSystem, + stateB: ChemicalSystem, + mapping: ComponentMapping | list[ComponentMapping] | None, + extends: ProtocolDAGResult | None, + ) -> list[ProtocolUnit]: + if extends: # pragma: no cover + raise NotImplementedError("Can't extend simulations yet") + + # Get alchemical components & validate them + mapping + alchem_comps = system_validation.get_alchemical_components(stateA, stateB) + _validate_alchemical_components(alchem_comps, mapping) + ligandmapping = mapping[0] if isinstance(mapping, list) else mapping + + # Validate solvent component + nonbond = self.settings.forcefield_settings.nonbonded_method + system_validation.validate_solvent(stateA, nonbond) + + # Validate protein component + system_validation.validate_protein(stateA) + + # actually create and return Units + Anames = ",".join(c.name for c in alchem_comps["stateA"]) + Bnames = ",".join(c.name for c in alchem_comps["stateB"]) + # our DAG has no dependencies, so just list units + n_repeats = self.settings.protocol_repeats + units = [ + HybridTopProtocolUnit( + protocol=self, + stateA=stateA, + stateB=stateB, + ligandmapping=ligandmapping, + generation=0, + repeat_id=int(uuid.uuid4()), + name=f"{Anames} to {Bnames} repeat {i} generation 0", + ) + for i in range(n_repeats) + ] + + return units diff --git a/src/pontibus/protocols/relative/hybridtop_units.py b/src/pontibus/protocols/relative/hybridtop_units.py new file mode 100644 index 0000000..8a12b6c --- /dev/null +++ b/src/pontibus/protocols/relative/hybridtop_units.py @@ -0,0 +1,677 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe +""" +ProtocolUnit implementations for the HybridTopProtocol. +""" + +import logging +import os +import pathlib +import warnings +from itertools import chain +from typing import Any + +import mdtraj +import numpy as np +import openmmtools +from gufe import SmallMoleculeComponent, SolventComponent +from gufe.settings import ThermoSettings +from gufe.vendor.openff.models.types import ArrayQuantity +from openfe.protocols.openmm_rfe import _rfe_utils +from openfe.protocols.openmm_rfe.equil_rfe_methods import ( + RelativeHybridTopologyProtocolUnit, + _get_alchemical_charge_difference, +) +from openfe.protocols.openmm_rfe.equil_rfe_settings import ( + AlchemicalSettings, + LambdaSettings, +) +from openfe.protocols.openmm_utils import ( + multistate_analysis, + omm_compute, + settings_validation, + system_validation, +) +from openfe.protocols.openmm_utils.omm_settings import ( + BasePartialChargeSettings, + IntegratorSettings, + MultiStateOutputSettings, + MultiStateSimulationSettings, +) +from openfe.utils import without_oechem_backend +from openff.interchange import Interchange +from openff.interchange.interop.openmm import to_openmm_positions +from openff.toolkit import Molecule as OFFMolecule +from openff.units import Quantity, unit +from openff.units.openmm import from_openmm, to_openmm +from openmm import CMMotionRemover, MonteCarloBarostat, System +from openmm import unit as omm_unit +from openmm.app import Topology +from openmmtools import multistate + +from pontibus.protocols.relative.settings import HybridTopProtocolSettings +from pontibus.protocols.solvation.base import _get_and_charge_solvent_offmol +from pontibus.utils.settings import ( + InterchangeFFSettings, + PackmolSolvationSettings, +) +from pontibus.utils.system_creation import ( + _get_force_field, + interchange_packmol_creation, +) +from pontibus.utils.system_manipulation import ( + adjust_system, + copy_interchange_with_replacement, +) + +logger = logging.getLogger(__name__) + + +class HybridTopProtocolUnit(RelativeHybridTopologyProtocolUnit): + @staticmethod + def _check_position_overlap( + mapping: dict[str, dict[int, int]], + positionsA: ArrayQuantity, + positionsB: ArrayQuantity, + threshold: Quantity = 1.0 * unit.angstrom, + ): + """ + Sanity check the overlap in positions. + + Parameters + ---------- + mapping : dict[str, dict[int, int]] + The system mappings between the two sets of positions. + positionsA : openff.units.Quantity + The system A positions. + positionsB : openff.units.Quantity + The system B positions. + tolerance : openff.units.Quantity + The maximum deivation allowed before an error or warning is raised. + + Raises + ------ + ValueError + If any env atoms deviate by more than the threshold. + UserWarning + If any core atoms deviate by more than the threshold. + """ + # Check env mappings + for key, val in mapping["old_to_new_env_atom_map"].items(): + if np.any(np.abs(positionsB[val] - positionsA[key]) > threshold): + msg = f"env mapping {key} : {val} deviates by more than {threshold}" + raise ValueError(msg) + + # Check core mappings + for key, val in mapping["old_to_new_core_atom_map"].items(): + if np.any(np.abs(positionsB[val] - positionsA[key]) > threshold): + msg = f"core mapping {key} : {val} deviates by more than {threshold}" + warnings.warn(msg) + logging.warning(msg) + + @staticmethod + def _get_barostat( + solvent_component: SolventComponent | None, + thermo_settings: ThermoSettings, + integrator_settings: IntegratorSettings, + ) -> MonteCarloBarostat | None: + """ + Helper to get a barostat for the system. + + Parameters + ---------- + solvent_component: SolventComponent | None + The system's SolventComponent, if there is one. + thermo_settings : ThermoSettings + The thermodynamic settings. + integrator_settings : IntegratorSettings + The integrator settings + + Returns + ------- + MonteCarloBarostat | None + None is there is no solvent, a MonteCarloBarostat otherwise. + """ + if solvent_component is None: + return None + + return MonteCarloBarostat( + to_openmm(thermo_settings.pressure), + to_openmm(thermo_settings.temperature), + integrator_settings.barostat_frequency.m, + ) + + @staticmethod + def _get_interchanges( + small_mols, + protein_component, + solvent_component, + forcefield_settings, + solvation_settings, + charge_settings, + ): + # Create an smc comp dictionary for stateA + stateA_smc_comps = dict(chain(small_mols["stateA"], small_mols["both"])) + + # Get solvent offmol if necessary + if solvent_component is not None: + solvent_offmol = _get_and_charge_solvent_offmol( + solvent_component, + solvation_settings, + charge_settings, + ) + else: + solvent_offmol = None + + # Get the stateA interchange + with without_oechem_backend(): + interA, comp_residsA = interchange_packmol_creation( + ffsettings=forcefield_settings, + solvation_settings=solvation_settings, + smc_components=stateA_smc_comps, + protein_component=protein_component, + solvent_component=solvent_component, + solvent_offmol=solvent_offmol, + ) + + # Get a list of the charged molecules to create stateB + stateB_charged_mols = [pair[1] for pair in chain(small_mols["stateB"], small_mols["both"])] + if solvent_component is not None and solvation_settings.assign_solvent_charges: + stateB_charged_mols.append(solvent_offmol) + + # Set the stateB interchange + interB = copy_interchange_with_replacement( + interchange=interA, + del_mol=small_mols["stateA"][0][1], + insert_mol=small_mols["stateB"][0][1], + force_field=_get_force_field(forcefield_settings), + charged_molecules=stateB_charged_mols, + ) + + # Fetch the alchemical resids for each state from the comp_resids + alchem_resids = { + "stateA": comp_residsA[small_mols["stateA"][0][0]], + "stateB": np.array([interB.topology.n_molecules - 1], dtype=int), + } + + return interA, interB, alchem_resids + + def _get_omm_objects( + self, + interchange: Interchange, + forcefield_settings: InterchangeFFSettings, + thermo_settings: ThermoSettings, + integrator_settings: IntegratorSettings, + solvent_component: SolventComponent | None, + ) -> tuple[Topology, omm_unit.Quantity, System]: + """ + Helper method to extract OpenMM objects from an Interchange object. + + Parameters + ---------- + interchange : Interchange + The Interchange object to get OpenMM objects from. + forcefield_settings : InterchangeFFSettings + The force field settings + thermo_settings : ThermoSettings + The thermodynamic parameter settings. + integrator_settings : IntegratorSettings + The integrator settings. + solvent_component : SolventComponent | None + The SolventComponent, if there is one. + """ + topology = interchange.to_openmm_topology(collate=True) + positions = to_openmm_positions( + interchange, + include_virtual_sites=True, + ) + system = interchange.to_openmm_system(hydrogen_mass=forcefield_settings.hydrogen_mass) + adjust_system( + system=system, + remove_force_types=CMMotionRemover, + add_forces=self._get_barostat( + solvent_component=solvent_component, + thermo_settings=thermo_settings, + integrator_settings=integrator_settings, + ), + ) + return topology, positions, system + + def run( + self, *, dry=False, verbose=True, scratch_basepath=None, shared_basepath=None + ) -> dict[str, Any]: + """Run the relative free energy calculation. + + Parameters + ---------- + dry : bool + Do a dry run of the calculation, creating all necessary hybrid + system components (topology, system, sampler, etc...) but without + running the simulation. + verbose : bool + Verbose output of the simulation progress. Output is provided via + INFO level logging. + scratch_basepath: Pathlike, optional + Where to store temporary files, defaults to current working directory + shared_basepath : Pathlike, optional + Where to run the calculation, defaults to current working directory + + Returns + ------- + dict + Outputs created in the basepath directory or the debug objects + (i.e. sampler) if ``dry==True``. + + Raises + ------ + error + Exception if anything failed + """ + if verbose: + self.logger.info("Preparing the hybrid topology simulation") + if scratch_basepath is None: + scratch_basepath = pathlib.Path(".") + if shared_basepath is None: + # use cwd + shared_basepath = pathlib.Path(".") + + # 0. General setup and settings dependency resolution step + + # Extract relevant settings + protocol_settings: HybridTopProtocolSettings = self._inputs["protocol"].settings + stateA = self._inputs["stateA"] + stateB = self._inputs["stateB"] + mapping = self._inputs["ligandmapping"] + + forcefield_settings: InterchangeFFSettings = protocol_settings.forcefield_settings + thermo_settings: ThermoSettings = protocol_settings.thermo_settings + alchem_settings: AlchemicalSettings = protocol_settings.alchemical_settings + lambda_settings: LambdaSettings = protocol_settings.lambda_settings + charge_settings: BasePartialChargeSettings = protocol_settings.partial_charge_settings + solvation_settings: PackmolSolvationSettings = protocol_settings.solvation_settings + sampler_settings: MultiStateSimulationSettings = protocol_settings.simulation_settings + output_settings: MultiStateOutputSettings = protocol_settings.output_settings + integrator_settings: IntegratorSettings = protocol_settings.integrator_settings + + # is the timestep good for the mass? + settings_validation.validate_timestep( + forcefield_settings.hydrogen_mass, integrator_settings.timestep + ) + # TODO: Also validate various conversions? + # Convert various time based inputs to steps/iterations + steps_per_iteration = settings_validation.convert_steps_per_iteration( + simulation_settings=sampler_settings, + integrator_settings=integrator_settings, + ) + + equil_steps = settings_validation.get_simsteps( + sim_length=sampler_settings.equilibration_length, + timestep=integrator_settings.timestep, + mc_steps=steps_per_iteration, + ) + prod_steps = settings_validation.get_simsteps( + sim_length=sampler_settings.production_length, + timestep=integrator_settings.timestep, + mc_steps=steps_per_iteration, + ) + + solvent_comp, protein_comp, small_mols = system_validation.get_components(stateA) + alchem_comps = system_validation.get_alchemical_components(stateA, stateB) + # We already do this in the Protocol but check the number of alchemical comps + for state in alchem_comps.values(): + assert len(state) == 1, "too many alchemical components found in one state" + + # Get the change difference between the end states + # and check if the charge correction used is appropriate + charge_difference = _get_alchemical_charge_difference( + mapping, + forcefield_settings.nonbonded_method, + alchem_settings.explicit_charge_correction, + solvent_comp, + ) + + # 1. Create stateA system + self.logger.info("Parameterizing molecules") + + # a. create (SMC, offmol) dictionaries and assign partial charges + # calculate partial charges manually if not already given + # convert to OpenFF here, + # and keep the molecule around to maintain the partial charges + off_small_mols: dict[str, list[tuple[SmallMoleculeComponent, OFFMolecule]]] + off_small_mols = { + "stateA": [(alchem_comps["stateA"][0], alchem_comps["stateA"][0].to_openff())], + "stateB": [(alchem_comps["stateB"][0], alchem_comps["stateB"][0].to_openff())], + "both": [ + (m, m.to_openff()) + for m in small_mols + if (m != alchem_comps["stateA"][0] and m != alchem_comps["stateB"][0]) + ], + } + + self._assign_partial_charges(charge_settings, off_small_mols) + + # Get stateA and stateB interchanges + stateA_interchange, stateB_interchange, alchem_resids = self._get_interchanges( + off_small_mols, + protein_comp, + solvent_comp, + forcefield_settings, + solvation_settings, + charge_settings, + ) + + # get topology & positions + stateA_topology, stateA_positions, stateA_system = self._get_omm_objects( + interchange=stateA_interchange, + forcefield_settings=forcefield_settings, + thermo_settings=thermo_settings, + integrator_settings=integrator_settings, + solvent_component=solvent_comp, + ) + stateB_topology, stateB_positions, stateB_system = self._get_omm_objects( + interchange=stateB_interchange, + forcefield_settings=forcefield_settings, + thermo_settings=thermo_settings, + integrator_settings=integrator_settings, + solvent_component=solvent_comp, + ) + + # c. Define correspondence mappings between the two systems + ligand_mappings = _rfe_utils.topologyhelpers.get_system_mappings( + mapping.componentA_to_componentB, + stateA_system, + stateA_topology, + alchem_resids["stateA"], + stateB_system, + stateB_topology, + alchem_resids["stateB"], + # These are non-optional settings for this method + fix_constraints=True, + ) + + # Sanity check the mappings looking at position overlaps + self._check_position_overlap( + ligand_mappings, + from_openmm(stateA_positions), + from_openmm(stateB_positions), + ) + + # d. if a charge correction is necessary, select alchemical waters + # and transform them + if alchem_settings.explicit_charge_correction: + alchem_water_resids = _rfe_utils.topologyhelpers.get_alchemical_waters( + stateA_topology, + from_openmm(stateA_positions).m, + charge_difference, + alchem_settings.explicit_charge_correction_cutoff, + ) + _rfe_utils.topologyhelpers.handle_alchemical_waters( + alchem_water_resids, + stateB_topology, + stateB_system, + ligand_mappings, + charge_difference, + solvent_comp, + ) + + # 3. Create the hybrid topology + hybrid_factory = _rfe_utils.relative.HybridTopologyFactory( + stateA_system, + stateA_positions, + stateA_topology, + stateB_system, + stateB_positions, + stateB_topology, + old_to_new_atom_map=ligand_mappings["old_to_new_atom_map"], + old_to_new_core_atom_map=ligand_mappings["old_to_new_core_atom_map"], + use_dispersion_correction=alchem_settings.use_dispersion_correction, + softcore_alpha=alchem_settings.softcore_alpha, + softcore_LJ_v2=alchem_settings.softcore_LJ.lower() == "gapsys", + softcore_LJ_v2_alpha=alchem_settings.softcore_alpha, + interpolate_old_and_new_14s=alchem_settings.turn_off_core_unique_exceptions, + ) + + # 4. Create lambda schedule + lambdas = _rfe_utils.lambdaprotocol.LambdaProtocol( + functions=lambda_settings.lambda_functions, windows=lambda_settings.lambda_windows + ) + + # pin lambda schedule spacing to n_replicas + n_replicas = sampler_settings.n_replicas + if n_replicas != len(lambdas.lambda_schedule): + errmsg = ( + f"Number of replicas {n_replicas} " + f"does not equal the number of lambda windows " + f"{len(lambdas.lambda_schedule)}" + ) + raise ValueError(errmsg) + + # 9. Create the multistate reporter + # Get the sub selection of the system to print coords for + selection_indices = hybrid_factory.hybrid_topology.select(output_settings.output_indices) + + # a. Create the multistate reporter + # convert checkpoint_interval from time to iterations + chk_intervals = settings_validation.convert_checkpoint_interval_to_iterations( + checkpoint_interval=output_settings.checkpoint_interval, + time_per_iteration=sampler_settings.time_per_iteration, + ) + + nc = shared_basepath / output_settings.output_filename + chk = output_settings.checkpoint_storage_filename + + if output_settings.positions_write_frequency is not None: + pos_interval = settings_validation.divmod_time_and_check( + numerator=output_settings.positions_write_frequency, + denominator=sampler_settings.time_per_iteration, + numerator_name="output settings' position_write_frequency", + denominator_name="sampler settings' time_per_iteration", + ) + else: + pos_interval = 0 + + if output_settings.velocities_write_frequency is not None: + vel_interval = settings_validation.divmod_time_and_check( + numerator=output_settings.velocities_write_frequency, + denominator=sampler_settings.time_per_iteration, + numerator_name="output settings' velocity_write_frequency", + denominator_name="sampler settings' time_per_iteration", + ) + else: + vel_interval = 0 + + reporter = multistate.MultiStateReporter( + storage=nc, + analysis_particle_indices=selection_indices, + checkpoint_interval=chk_intervals, + checkpoint_storage=chk, + position_interval=pos_interval, + velocity_interval=vel_interval, + ) + + # b. Write out a PDB containing the subsampled hybrid state + bfactors = np.zeros_like(selection_indices, dtype=float) # solvent + bfactors[ + np.in1d(selection_indices, list(hybrid_factory._atom_classes["unique_old_atoms"])) + ] = 0.25 # lig A + bfactors[np.in1d(selection_indices, list(hybrid_factory._atom_classes["core_atoms"]))] = ( + 0.50 # core + ) + bfactors[ + np.in1d(selection_indices, list(hybrid_factory._atom_classes["unique_new_atoms"])) + ] = 0.75 # lig B + + if len(selection_indices) > 0: + traj = mdtraj.Trajectory( + hybrid_factory.hybrid_positions[selection_indices, :], + hybrid_factory.hybrid_topology.subset(selection_indices), + ) + traj.save_pdb( + shared_basepath / output_settings.output_structure, + bfactors=bfactors, + ) + + # 10. Get compute platform + # restrict to a single CPU if running vacuum + restrict_cpu = forcefield_settings.nonbonded_method.lower() == "nocutoff" + platform = omm_compute.get_openmm_platform( + platform_name=protocol_settings.engine_settings.compute_platform, + gpu_device_index=protocol_settings.engine_settings.gpu_device_index, + restrict_cpu_count=restrict_cpu, + ) + + # 11. Set the integrator + # a. Validate integrator settings for current system + # Virtual sites sanity check - ensure we restart velocities when + # there are virtual sites in the system + if hybrid_factory.has_virtual_sites: + if not integrator_settings.reassign_velocities: + errmsg = ( + "Simulations with virtual sites without velocity " + "reassignments are unstable in openmmtools" + ) + raise ValueError(errmsg) + + # b. create langevin integrator + integrator = openmmtools.mcmc.LangevinDynamicsMove( + timestep=to_openmm(integrator_settings.timestep), + collision_rate=to_openmm(integrator_settings.langevin_collision_rate), + n_steps=steps_per_iteration, + reassign_velocities=integrator_settings.reassign_velocities, + n_restart_attempts=integrator_settings.n_restart_attempts, + constraint_tolerance=integrator_settings.constraint_tolerance, + ) + + # 12. Create sampler + self.logger.info("Creating and setting up the sampler") + rta_its, rta_min_its = settings_validation.convert_real_time_analysis_iterations( + simulation_settings=sampler_settings, + ) + # convert early_termination_target_error from kcal/mol to kT + early_termination_target_error = ( + settings_validation.convert_target_error_from_kcal_per_mole_to_kT( + thermo_settings.temperature, + sampler_settings.early_termination_target_error, + ) + ) + + if sampler_settings.sampler_method.lower() == "repex": + sampler = _rfe_utils.multistate.HybridRepexSampler( + mcmc_moves=integrator, + hybrid_factory=hybrid_factory, + online_analysis_interval=rta_its, + online_analysis_target_error=early_termination_target_error, + online_analysis_minimum_iterations=rta_min_its, + ) + elif sampler_settings.sampler_method.lower() == "sams": + sampler = _rfe_utils.multistate.HybridSAMSSampler( + mcmc_moves=integrator, + hybrid_factory=hybrid_factory, + online_analysis_interval=rta_its, + online_analysis_minimum_iterations=rta_min_its, + flatness_criteria=sampler_settings.sams_flatness_criteria, + gamma0=sampler_settings.sams_gamma0, + ) + elif sampler_settings.sampler_method.lower() == "independent": + sampler = _rfe_utils.multistate.HybridMultiStateSampler( + mcmc_moves=integrator, + hybrid_factory=hybrid_factory, + online_analysis_interval=rta_its, + online_analysis_target_error=early_termination_target_error, + online_analysis_minimum_iterations=rta_min_its, + ) + + else: + raise AttributeError(f"Unknown sampler {sampler_settings.sampler_method}") + + sampler.setup( + n_replicas=sampler_settings.n_replicas, + reporter=reporter, + lambda_protocol=lambdas, + temperature=to_openmm(thermo_settings.temperature), + endstates=alchem_settings.endstate_dispersion_correction, + minimization_platform=platform.getName(), + ) + + try: + # Create context caches (energy + sampler) + energy_context_cache = openmmtools.cache.ContextCache( + capacity=None, + time_to_live=None, + platform=platform, + ) + + sampler_context_cache = openmmtools.cache.ContextCache( + capacity=None, + time_to_live=None, + platform=platform, + ) + + sampler.energy_context_cache = energy_context_cache + sampler.sampler_context_cache = sampler_context_cache + + if not dry: # pragma: no-cover + # minimize + if verbose: + self.logger.info("Running minimization") + + sampler.minimize(max_iterations=sampler_settings.minimization_steps) + + # equilibrate + if verbose: + self.logger.info("Running equilibration phase") + + sampler.equilibrate(int(equil_steps / steps_per_iteration)) + + # production + if verbose: + self.logger.info("Running production phase") + + sampler.extend(int(prod_steps / steps_per_iteration)) + + self.logger.info("Production phase complete") + + self.logger.info("Post-simulation analysis of results") + # calculate relevant analyses of the free energies & sampling + # First close & reload the reporter to avoid netcdf clashes + analyzer = multistate_analysis.MultistateEquilFEAnalysis( + reporter, + sampling_method=sampler_settings.sampler_method.lower(), + result_units=unit.kilocalorie_per_mole, + ) + analyzer.plot(filepath=shared_basepath, filename_prefix="") + analyzer.close() + + else: + # clean up the reporter file + fns = [ + shared_basepath / output_settings.output_filename, + shared_basepath / output_settings.checkpoint_storage_filename, + ] + for fn in fns: + os.remove(fn) + finally: + # close reporter when you're done, prevent + # file handle clashes + reporter.close() + + # clear GPU contexts + # TODO: use cache.empty() calls when openmmtools #690 is resolved + # replace with above + for context in list(energy_context_cache._lru._data.keys()): + del energy_context_cache._lru._data[context] + for context in list(sampler_context_cache._lru._data.keys()): + del sampler_context_cache._lru._data[context] + # cautiously clear out the global context cache too + for context in list(openmmtools.cache.global_context_cache._lru._data.keys()): + del openmmtools.cache.global_context_cache._lru._data[context] + + del sampler_context_cache, energy_context_cache + + if not dry: + del integrator, sampler + + if not dry: # pragma: no-cover + return {"nc": nc, "last_checkpoint": chk, **analyzer.unit_results_dict} + else: + return {"debug": {"sampler": sampler}} diff --git a/src/pontibus/protocols/relative/settings.py b/src/pontibus/protocols/relative/settings.py new file mode 100644 index 0000000..2775f99 --- /dev/null +++ b/src/pontibus/protocols/relative/settings.py @@ -0,0 +1,38 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe + +"""Settings classes for RFE Protocols using OpenMM + OpenMMTools + +This module implements the necessary settings necessary to run the following: + * HybridTopProtocol + +See Also +-------- +pontibus.protocols.relative.HybridTopProtocol +""" + +from openfe.protocols.openmm_rfe.equil_rfe_settings import ( + RelativeHybridTopologyProtocolSettings, +) + +from pontibus.utils.settings import ( + InterchangeFFSettings, + PackmolSolvationSettings, +) + + +class HybridTopProtocolSettings(RelativeHybridTopologyProtocolSettings): + """ + Configuration object for ``HybridTopologyRFEProtocol``. + + See Also + -------- + pontibus.protocols.relative.HybridTopologyRFEProtocol + """ + + # Inherited things + forcefield_settings: InterchangeFFSettings + """Parameters to set up in the force field""" + + solvation_settings: PackmolSolvationSettings + """Settings for solvating the system.""" diff --git a/src/pontibus/protocols/solvation/__init__.py b/src/pontibus/protocols/solvation/__init__.py index b98b8ac..541b58e 100644 --- a/src/pontibus/protocols/solvation/__init__.py +++ b/src/pontibus/protocols/solvation/__init__.py @@ -2,7 +2,6 @@ # For details, see https://github.com/OpenFreeEnergy/openfe """ Run absolute solvation free energy calculations using OpenMM and OpenMMTools. - """ from .asfe_protocol import ( diff --git a/src/pontibus/protocols/solvation/base.py b/src/pontibus/protocols/solvation/base.py index 5a4d8e5..4f937a8 100644 --- a/src/pontibus/protocols/solvation/base.py +++ b/src/pontibus/protocols/solvation/base.py @@ -15,10 +15,8 @@ ) from gufe.settings import SettingsBaseModel from openfe.protocols.openmm_afe.base import BaseAbsoluteUnit -from openfe.protocols.openmm_utils import charge_generation from openfe.protocols.openmm_utils.omm_settings import ( IntegratorSettings, - OpenFFPartialChargeSettings, ) from openfe.utils import log_system_probe, without_oechem_backend from openff.interchange.interop.openmm import to_openmm_positions @@ -30,12 +28,12 @@ AlchemicalRegion, ) -from pontibus.components import ExtendedSolventComponent -from pontibus.protocols.solvation.settings import PackmolSolvationSettings from pontibus.utils.experimental_absolute_factory import ( ExperimentalAbsoluteAlchemicalFactory, ) +from pontibus.utils.protocol_utils import _get_and_charge_solvent_offmol from pontibus.utils.system_creation import interchange_packmol_creation +from pontibus.utils.system_manipulation import adjust_system logger = logging.getLogger(__name__) @@ -43,58 +41,6 @@ class BaseASFEUnit(BaseAbsoluteUnit): _simtype: str - @staticmethod - def _get_and_charge_solvent_offmol( - solvent_component: SolventComponent | ExtendedSolventComponent, - solvation_settings: PackmolSolvationSettings, - partial_charge_settings: OpenFFPartialChargeSettings, - ) -> OFFMolecule: - """ - Helper method to fetch the solvent offmol either - from an existing solvent_smcs, or from smiles. - - Parameters - ---------- - solvent_component : SolventComponent - smiles for the solvent molecule - solvation_settings : PackmolSolvationSettings - Settings defining how the system will be solvated - partial_charge_settings : OpenFFPartialChargeSettigns - Settings defining how partial charges are applied - - Returns - ------- - offmol : openff.toolkit.Molecule - - Notes - ----- - * If created from a smiles, the solvent will be assigned - a single conformer through `Molecule.generate_conformers`. - """ - # Get the solvent offmol - if isinstance(solvent_component, ExtendedSolventComponent): - solvent_offmol = solvent_component.solvent_molecule.to_openff() # type: ignore[union-attr] - else: - # If not, we create the solvent from smiles - # We generate a single conformer to avoid packing issues - solvent_offmol = OFFMolecule.from_smiles(solvent_component.smiles) - solvent_offmol.generate_conformers(n_conformers=1) - - # In-place assign solvent offmol charges if necessary - # Note: we don't enforce partial charge assignment to avoid - # cases where we want to rely on library charges instead. - if solvation_settings.assign_solvent_charges: - charge_generation.assign_offmol_partial_charges( - offmol=solvent_offmol, - overwrite=False, - method=partial_charge_settings.partial_charge_method, - toolkit_backend=partial_charge_settings.off_toolkit_backend, - generate_n_conformers=partial_charge_settings.number_of_conformers, - nagl_model=partial_charge_settings.nagl_model, - ) - - return solvent_offmol - @staticmethod def _validate_vsites(system: openmm.System, integrator_settings: IntegratorSettings) -> None: """ @@ -173,7 +119,7 @@ def _get_omm_objects( # Get solvent offmol if necessary if solvent_component is not None: - solvent_offmol = self._get_and_charge_solvent_offmol( + solvent_offmol = _get_and_charge_solvent_offmol( solvent_component, settings["solvation_settings"], settings["charge_settings"], @@ -198,13 +144,6 @@ def _get_omm_objects( hydrogen_mass=settings["forcefield_settings"].hydrogen_mass ) - # Pull out the CMMotionRemover - # TODO: add test that checks the number of forces - for idx in reversed(range(omm_system.getNumForces())): - force = omm_system.getForce(idx) - if isinstance(force, openmm.CMMotionRemover): - omm_system.removeForce(idx) - # Add a barostat if needed if solvent_component is not None: barostat = openmm.MonteCarloBarostat( @@ -212,7 +151,12 @@ def _get_omm_objects( to_openmm(settings["thermo_settings"].temperature), settings["integrator_settings"].barostat_frequency.m, ) - omm_system.addForce(barostat) + else: + barostat = None + + adjust_system( + system=omm_system, remove_force_types=openmm.CMMotionRemover, add_forces=barostat + ) positions = to_openmm_positions(interchange, include_virtual_sites=True) diff --git a/src/pontibus/tests/conftest.py b/src/pontibus/tests/conftest.py index 8a5c8b5..13ab201 100644 --- a/src/pontibus/tests/conftest.py +++ b/src/pontibus/tests/conftest.py @@ -6,8 +6,10 @@ from importlib import resources import gufe +import openfe import pytest from gufe import SmallMoleculeComponent +from openff.units import unit from rdkit import Chem @@ -138,7 +140,7 @@ def pytest_configure(config): @pytest.fixture(scope="session") def benzene_modifications(): files = {} - with importlib.resources.files("openfe.tests.data") as d: + with importlib.resources.files("pontibus.tests.data") as d: fn = str(d / "benzene_modifications.sdf") supp = Chem.SDMolSupplier(str(fn), removeHs=False) for rdmol in supp: @@ -146,6 +148,17 @@ def benzene_modifications(): return files +@pytest.fixture(scope="session") +def benzene_modifications_charged(): + files = {} + with importlib.resources.files("pontibus.tests.data") as d: + fn = str(d / "benzene_modifications_charged.sdf") + supp = Chem.SDMolSupplier(str(fn), removeHs=False) + for rdmol in supp: + files[rdmol.GetProp("_Name")] = SmallMoleculeComponent(rdmol) + return files + + @pytest.fixture() def CN_molecule(): """ @@ -167,3 +180,77 @@ def T4_protein_component(): comp = gufe.ProteinComponent.from_pdb_file(fn, name="T4_protein") return comp + + +@pytest.fixture +def benzene_vacuum_system(benzene_modifications_charged): + return openfe.ChemicalSystem( + {"ligand": benzene_modifications_charged["benzene"]}, + ) + + +@pytest.fixture(scope="session") +def benzene_system(benzene_modifications_charged): + return openfe.ChemicalSystem( + { + "ligand": benzene_modifications_charged["benzene"], + "solvent": openfe.SolventComponent( + positive_ion="Na", negative_ion="Cl", ion_concentration=0.15 * unit.molar + ), + }, + ) + + +@pytest.fixture +def benzene_complex_system(benzene_modifications_charged, T4_protein_component): + return openfe.ChemicalSystem( + { + "ligand": benzene_modifications_charged["benzene"], + "solvent": openfe.SolventComponent( + positive_ion="Na", negative_ion="Cl", ion_concentration=0.15 * unit.molar + ), + "protein": T4_protein_component, + } + ) + + +@pytest.fixture +def toluene_vacuum_system(benzene_modifications_charged): + return openfe.ChemicalSystem( + {"ligand": benzene_modifications_charged["toluene"]}, + ) + + +@pytest.fixture(scope="session") +def toluene_system(benzene_modifications_charged): + return openfe.ChemicalSystem( + { + "ligand": benzene_modifications_charged["toluene"], + "solvent": openfe.SolventComponent( + positive_ion="Na", negative_ion="Cl", ion_concentration=0.15 * unit.molar + ), + }, + ) + + +@pytest.fixture +def toluene_complex_system(benzene_modifications_charged, T4_protein_component): + return openfe.ChemicalSystem( + { + "ligand": benzene_modifications_charged["toluene"], + "solvent": openfe.SolventComponent( + positive_ion="Na", negative_ion="Cl", ion_concentration=0.15 * unit.molar + ), + "protein": T4_protein_component, + } + ) + + +@pytest.fixture(scope="session") +def benzene_to_toluene_mapping(benzene_modifications_charged): + mapper = openfe.setup.LomapAtomMapper(element_change=False) + + molA = benzene_modifications_charged["benzene"] + molB = benzene_modifications_charged["toluene"] + + return next(mapper.suggest_mappings(molA, molB)) diff --git a/src/pontibus/tests/data/benzene_modifications_charged.sdf b/src/pontibus/tests/data/benzene_modifications_charged.sdf new file mode 100644 index 0000000..f1cebed --- /dev/null +++ b/src/pontibus/tests/data/benzene_modifications_charged.sdf @@ -0,0 +1,295 @@ +benzene + RDKit 3D + + 12 12 0 0 0 0 0 0 0 0999 V2000 + 25.9780 5.3270 4.7790 C 0 0 0 0 0 0 0 0 0 0 0 0 + 26.3950 5.0740 3.4990 C 0 0 0 0 0 0 0 0 0 0 0 0 + 27.3400 5.8600 2.9020 C 0 0 0 0 0 0 0 0 0 0 0 0 + 27.8370 6.9210 3.5690 C 0 0 0 0 0 0 0 0 0 0 0 0 + 27.4200 7.1960 4.8560 C 0 0 0 0 0 0 0 0 0 0 0 0 + 26.4980 6.3790 5.4690 C 0 0 0 0 0 0 0 0 0 0 0 0 + 25.2298 4.6859 5.2451 H 0 0 0 0 0 0 0 0 0 0 0 0 + 25.9676 4.2351 2.9497 H 0 0 0 0 0 0 0 0 0 0 0 0 + 27.6890 5.6311 1.8951 H 0 0 0 0 0 0 0 0 0 0 0 0 + 28.5730 7.5660 3.0889 H 0 0 0 0 0 0 0 0 0 0 0 0 + 27.8209 8.0598 5.3863 H 0 0 0 0 0 0 0 0 0 0 0 0 + 26.1874 6.5720 6.4958 H 0 0 0 0 0 0 0 0 0 0 0 0 + 1 2 2 0 + 1 6 1 0 + 1 7 1 0 + 2 3 1 0 + 2 8 1 0 + 3 4 2 0 + 3 9 1 0 + 4 5 1 0 + 4 10 1 0 + 5 6 2 0 + 5 11 1 0 + 6 12 1 0 +M END + +> +benzene + +> +-0.13 -0.13 -0.13 -0.13 -0.13 -0.13 0.13 0.13 0.13 0.13 0.13 0.13 + +$$$$ +phenol + RDKit 3D + + 13 13 0 0 0 0 0 0 0 0999 V2000 + 25.9780 5.3270 4.7790 C 0 0 0 0 0 0 0 0 0 0 0 0 + 26.3950 5.0740 3.4990 C 0 0 0 0 0 0 0 0 0 0 0 0 + 27.3400 5.8600 2.9020 C 0 0 0 0 0 0 0 0 0 0 0 0 + 27.8370 6.9210 3.5690 C 0 0 0 0 0 0 0 0 0 0 0 0 + 27.4200 7.1960 4.8560 C 0 0 0 0 0 0 0 0 0 0 0 0 + 26.4980 6.3790 5.4690 C 0 0 0 0 0 0 0 0 0 0 0 0 + 25.2298 4.6859 5.2451 H 0 0 0 0 0 0 0 0 0 0 0 0 + 25.9676 4.2351 2.9497 H 0 0 0 0 0 0 0 0 0 0 0 0 + 27.6890 5.6311 1.8951 H 0 0 0 0 0 0 0 0 0 0 0 0 + 28.5730 7.5660 3.0889 H 0 0 0 0 0 0 0 0 0 0 0 0 + 28.1311 8.0887 6.4624 H 0 0 0 0 0 0 0 0 0 0 0 0 + 26.1874 6.5720 6.4958 H 0 0 0 0 0 0 0 0 0 0 0 0 + 27.9460 8.3293 5.5517 O 0 0 0 0 0 0 0 0 0 0 0 0 + 1 2 2 0 + 1 6 1 0 + 1 7 1 0 + 2 3 1 0 + 2 8 1 0 + 3 4 2 0 + 3 9 1 0 + 4 5 1 0 + 4 10 1 0 + 5 6 2 0 + 5 13 1 0 + 6 12 1 0 + 11 13 1 0 +M END + +> +phenol + +> +-0.094423076923076915 -0.16592307692307692 -0.094423076923076915 -0.18492307692307691 0.12317692307692309 -0.18492307692307691 0.13307692307692309 0.13307692307692309 0.13307692307692309 +0.14157692307692307 0.41807692307692307 0.14157692307692307 -0.4990230769230769 + +$$$$ +benzonitrile + RDKit 3D + + 13 13 0 0 0 0 0 0 0 0999 V2000 + 28.5559 9.5700 6.2831 N 0 0 0 0 0 0 0 0 0 0 0 0 + 27.9981 8.4043 5.5824 C 0 0 0 0 0 0 0 0 0 0 0 0 + 25.9780 5.3270 4.7790 C 0 0 0 0 0 0 0 0 0 0 0 0 + 26.3950 5.0740 3.4990 C 0 0 0 0 0 0 0 0 0 0 0 0 + 27.3400 5.8600 2.9020 C 0 0 0 0 0 0 0 0 0 0 0 0 + 27.8370 6.9210 3.5690 C 0 0 0 0 0 0 0 0 0 0 0 0 + 27.4200 7.1960 4.8560 C 0 0 0 0 0 0 0 0 0 0 0 0 + 26.4980 6.3790 5.4690 C 0 0 0 0 0 0 0 0 0 0 0 0 + 25.2298 4.6859 5.2451 H 0 0 0 0 0 0 0 0 0 0 0 0 + 25.9676 4.2351 2.9497 H 0 0 0 0 0 0 0 0 0 0 0 0 + 27.6890 5.6311 1.8951 H 0 0 0 0 0 0 0 0 0 0 0 0 + 28.5730 7.5660 3.0889 H 0 0 0 0 0 0 0 0 0 0 0 0 + 26.1874 6.5720 6.4958 H 0 0 0 0 0 0 0 0 0 0 0 0 + 1 2 3 0 + 3 4 2 0 + 3 8 1 0 + 3 9 1 0 + 4 5 1 0 + 4 10 1 0 + 5 6 2 0 + 5 11 1 0 + 6 7 1 0 + 6 12 1 0 + 2 7 1 0 + 7 8 2 0 + 8 13 1 0 +M END + +> +benzonitrile + +> +-0.36380000000000001 0.23380000000000001 -0.13500000000000001 -0.107 -0.13500000000000001 -0.090999999999999998 -0.019000000000000003 -0.090999999999999998 0.14000000000000001 +0.13800000000000001 0.14000000000000001 0.14499999999999999 0.14499999999999999 + +$$$$ +benzaldehyde + RDKit 3D + + 14 14 0 0 0 0 0 0 0 0999 V2000 + 29.2079 8.8492 4.9632 O 0 0 0 0 0 0 0 0 0 0 0 0 + 27.5482 8.8691 6.4597 H 0 0 0 0 0 0 0 0 0 0 0 0 + 27.9981 8.4043 5.5824 C 0 0 0 0 0 0 0 0 0 0 0 0 + 25.9780 5.3270 4.7790 C 0 0 0 0 0 0 0 0 0 0 0 0 + 26.3950 5.0740 3.4990 C 0 0 0 0 0 0 0 0 0 0 0 0 + 27.3400 5.8600 2.9020 C 0 0 0 0 0 0 0 0 0 0 0 0 + 27.8370 6.9210 3.5690 C 0 0 0 0 0 0 0 0 0 0 0 0 + 27.4200 7.1960 4.8560 C 0 0 0 0 0 0 0 0 0 0 0 0 + 26.4980 6.3790 5.4690 C 0 0 0 0 0 0 0 0 0 0 0 0 + 25.2298 4.6859 5.2451 H 0 0 0 0 0 0 0 0 0 0 0 0 + 25.9676 4.2351 2.9497 H 0 0 0 0 0 0 0 0 0 0 0 0 + 27.6890 5.6311 1.8951 H 0 0 0 0 0 0 0 0 0 0 0 0 + 28.5730 7.5660 3.0889 H 0 0 0 0 0 0 0 0 0 0 0 0 + 26.1874 6.5720 6.4958 H 0 0 0 0 0 0 0 0 0 0 0 0 + 1 3 2 0 + 2 3 1 0 + 4 5 2 0 + 4 9 1 0 + 4 10 1 0 + 5 6 1 0 + 5 11 1 0 + 6 7 2 0 + 6 12 1 0 + 7 8 1 0 + 7 13 1 0 + 3 8 1 0 + 8 9 2 0 + 9 14 1 0 +M END + +> +benzaldehyde + +> +-0.52817142857142862 -0.0028714285714285795 0.5754285714285714 -0.14507142857142857 -0.098071428571428587 -0.14507142857142857 -0.078071428571428583 -0.19767142857142858 +-0.078071428571428583 0.13742857142857143 0.13492857142857143 0.13742857142857143 0.14392857142857141 0.14392857142857141 + +$$$$ +styrene + RDKit 3D + + 16 16 0 0 0 0 0 0 0 0999 V2000 + 29.2873 8.8784 4.9226 C 0 0 0 0 0 0 0 0 0 0 0 0 + 29.6609 8.3486 4.0463 H 0 0 0 0 0 0 0 0 0 0 0 0 + 29.8344 9.7353 5.3157 H 0 0 0 0 0 0 0 0 0 0 0 0 + 27.5365 8.8812 6.4825 H 0 0 0 0 0 0 0 0 0 0 0 0 + 27.9864 8.4164 5.6052 C 0 0 0 0 0 0 0 0 0 0 0 0 + 25.9780 5.3270 4.7790 C 0 0 0 0 0 0 0 0 0 0 0 0 + 26.3950 5.0740 3.4990 C 0 0 0 0 0 0 0 0 0 0 0 0 + 27.3400 5.8600 2.9020 C 0 0 0 0 0 0 0 0 0 0 0 0 + 27.8370 6.9210 3.5690 C 0 0 0 0 0 0 0 0 0 0 0 0 + 27.4200 7.1960 4.8560 C 0 0 0 0 0 0 0 0 0 0 0 0 + 26.4980 6.3790 5.4690 C 0 0 0 0 0 0 0 0 0 0 0 0 + 25.2298 4.6859 5.2451 H 0 0 0 0 0 0 0 0 0 0 0 0 + 25.9676 4.2351 2.9497 H 0 0 0 0 0 0 0 0 0 0 0 0 + 27.6890 5.6311 1.8951 H 0 0 0 0 0 0 0 0 0 0 0 0 + 28.5730 7.5660 3.0889 H 0 0 0 0 0 0 0 0 0 0 0 0 + 26.1874 6.5720 6.4958 H 0 0 0 0 0 0 0 0 0 0 0 0 + 1 2 1 0 + 1 3 1 0 + 1 5 2 0 + 4 5 1 0 + 6 7 2 0 + 6 11 1 0 + 6 12 1 0 + 7 8 1 0 + 7 13 1 0 + 8 9 2 0 + 8 14 1 0 + 9 10 1 0 + 9 15 1 0 + 5 10 1 0 + 10 11 2 0 + 11 16 1 0 +M END + +> +styrene + +> +-0.20899999999999999 0.11349999999999999 0.11349999999999999 0.123 -0.1132 -0.13100000000000001 -0.127 -0.13100000000000001 -0.11849999999999999 -0.057800000000000004 -0.11849999999999999 +0.13100000000000001 0.13100000000000001 0.13100000000000001 0.13150000000000001 0.13150000000000001 + +$$$$ +anisole + RDKit 3D + + 16 16 0 0 0 0 0 0 0 0999 V2000 + 29.2873 8.8784 4.9226 C 0 0 0 0 0 0 0 0 0 0 0 0 + 29.5502 9.7990 5.4437 H 0 0 0 0 0 0 0 0 0 0 0 0 + 28.0548 8.3459 5.4720 O 0 0 0 0 0 0 0 0 0 0 0 0 + 30.0866 8.1484 5.0502 H 0 0 0 0 0 0 0 0 0 0 0 0 + 29.1525 9.0868 3.8612 H 0 0 0 0 0 0 0 0 0 0 0 0 + 25.9780 5.3270 4.7790 C 0 0 0 0 0 0 0 0 0 0 0 0 + 26.3950 5.0740 3.4990 C 0 0 0 0 0 0 0 0 0 0 0 0 + 27.3400 5.8600 2.9020 C 0 0 0 0 0 0 0 0 0 0 0 0 + 27.8370 6.9210 3.5690 C 0 0 0 0 0 0 0 0 0 0 0 0 + 27.4200 7.1960 4.8560 C 0 0 0 0 0 0 0 0 0 0 0 0 + 26.4980 6.3790 5.4690 C 0 0 0 0 0 0 0 0 0 0 0 0 + 25.2298 4.6859 5.2451 H 0 0 0 0 0 0 0 0 0 0 0 0 + 25.9676 4.2351 2.9497 H 0 0 0 0 0 0 0 0 0 0 0 0 + 27.6890 5.6311 1.8951 H 0 0 0 0 0 0 0 0 0 0 0 0 + 28.5730 7.5660 3.0889 H 0 0 0 0 0 0 0 0 0 0 0 0 + 26.1874 6.5720 6.4958 H 0 0 0 0 0 0 0 0 0 0 0 0 + 1 2 1 0 + 1 4 1 0 + 1 5 1 0 + 1 3 1 0 + 6 7 2 0 + 6 11 1 0 + 6 12 1 0 + 7 8 1 0 + 7 13 1 0 + 8 9 2 0 + 8 14 1 0 + 9 10 1 0 + 9 15 1 0 + 3 10 1 0 + 10 11 2 0 + 11 16 1 0 +M END + +> +anisole + +> +0.113825 0.043825000000000003 -0.32877500000000004 0.043825000000000003 0.043825000000000003 -0.097875000000000004 -0.16487500000000002 -0.097875000000000004 -0.17987500000000001 0.123225 +-0.17987500000000001 0.13212499999999999 0.13212499999999999 0.13212499999999999 0.14212499999999997 0.14212499999999997 + +$$$$ +toluene + RDKit 3D + + 15 15 0 0 0 0 0 0 0 0999 V2000 + 28.9072 8.7434 5.1220 H 0 0 0 0 0 0 0 0 0 0 0 0 + 28.1966 8.1433 6.6393 H 0 0 0 0 0 0 0 0 0 0 0 0 + 27.9864 8.4164 5.6052 C 0 0 0 0 0 0 0 0 0 0 0 0 + 27.2579 9.2269 5.5838 H 0 0 0 0 0 0 0 0 0 0 0 0 + 25.9780 5.3270 4.7790 C 0 0 0 0 0 0 0 0 0 0 0 0 + 26.3950 5.0740 3.4990 C 0 0 0 0 0 0 0 0 0 0 0 0 + 27.3400 5.8600 2.9020 C 0 0 0 0 0 0 0 0 0 0 0 0 + 27.8370 6.9210 3.5690 C 0 0 0 0 0 0 0 0 0 0 0 0 + 27.4200 7.1960 4.8560 C 0 0 0 0 0 0 0 0 0 0 0 0 + 26.4980 6.3790 5.4690 C 0 0 0 0 0 0 0 0 0 0 0 0 + 25.2298 4.6859 5.2451 H 0 0 0 0 0 0 0 0 0 0 0 0 + 25.9676 4.2351 2.9497 H 0 0 0 0 0 0 0 0 0 0 0 0 + 27.6890 5.6311 1.8951 H 0 0 0 0 0 0 0 0 0 0 0 0 + 28.5730 7.5660 3.0889 H 0 0 0 0 0 0 0 0 0 0 0 0 + 26.1874 6.5720 6.4958 H 0 0 0 0 0 0 0 0 0 0 0 0 + 1 3 1 0 + 2 3 1 0 + 3 4 1 0 + 5 6 2 0 + 5 10 1 0 + 5 11 1 0 + 6 7 1 0 + 6 12 1 0 + 7 8 2 0 + 7 13 1 0 + 8 9 1 0 + 8 14 1 0 + 3 9 1 0 + 9 10 2 0 + 10 15 1 0 +M END + +> +toluene + +> +0.044033066666666669 0.044033066666666669 -0.053799933333333334 0.044033066666666669 -0.12699993333333334 -0.13499993333333335 -0.12699993333333334 -0.13099993333333335 +-0.077299933333333321 -0.13099993333333335 0.13000006666666666 0.13000006666666666 0.13000006666666666 0.13000006666666666 0.13000006666666666 + +$$$$ diff --git a/src/pontibus/tests/data/relative_protocol/HybridTopProtocol_solvent_json_results.gz b/src/pontibus/tests/data/relative_protocol/HybridTopProtocol_solvent_json_results.gz new file mode 100644 index 0000000..99ada77 Binary files /dev/null and b/src/pontibus/tests/data/relative_protocol/HybridTopProtocol_solvent_json_results.gz differ diff --git a/src/pontibus/tests/data/relative_protocol/HybridTopProtocol_vacuum_json_results.gz b/src/pontibus/tests/data/relative_protocol/HybridTopProtocol_vacuum_json_results.gz new file mode 100644 index 0000000..e08b42d Binary files /dev/null and b/src/pontibus/tests/data/relative_protocol/HybridTopProtocol_vacuum_json_results.gz differ diff --git a/src/pontibus/tests/data/relative_protocol/__init__.py b/src/pontibus/tests/data/relative_protocol/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/pontibus/tests/protocols/relative/__init__.py b/src/pontibus/tests/protocols/relative/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/pontibus/tests/protocols/relative/conftest.py b/src/pontibus/tests/protocols/relative/conftest.py new file mode 100644 index 0000000..12cf90d --- /dev/null +++ b/src/pontibus/tests/protocols/relative/conftest.py @@ -0,0 +1,32 @@ +import gzip +from importlib import resources + +import pytest + + +@pytest.fixture +def rfe_solv_transformation_json() -> str: + """ + HybridTop results object as created by quickrun. + + generated with devtools/gen-serialized-results.py + """ + d = resources.files("pontibus.tests.data.relative_protocol") + file = d / "HybridTopProtocol_solvent_json_results.gz" + + with gzip.open(file.as_posix(), "r") as f: # type: ignore + return f.read().decode() # type: ignore + + +@pytest.fixture +def rfe_vacuum_transformation_json() -> str: + """ + Hybrid results object as created by quickrun. + + generated with devtools/gen-serialized-results.py + """ + d = resources.files("pontibus.tests.data.relative_protocol") + file = d / "HybridTopProtocol_vacuum_json_results.gz" + + with gzip.open(file.as_posix(), "r") as f: # type: ignore + return f.read().decode() # type: ignore diff --git a/src/pontibus/tests/protocols/relative/test_protocol.py b/src/pontibus/tests/protocols/relative/test_protocol.py new file mode 100644 index 0000000..ff32afb --- /dev/null +++ b/src/pontibus/tests/protocols/relative/test_protocol.py @@ -0,0 +1,438 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe + +import mdtraj as mdt +import numpy as np +import openfe +import pytest +from openff.units import unit +from openff.units.openmm import ensure_quantity +from openmm import ( + CustomAngleForce, + CustomBondForce, + CustomNonbondedForce, + CustomTorsionForce, + HarmonicAngleForce, + HarmonicBondForce, + MonteCarloBarostat, + NonbondedForce, + PeriodicTorsionForce, +) +from openmm import unit as omm_unit +from openmmtools.multistate import MultiStateSampler +from rdkit import Chem + +from pontibus.protocols.relative import HybridTopProtocol, HybridTopProtocolUnit + + +def test_create_default_settings(): + settings = HybridTopProtocol.default_settings() + + assert settings + + +def test_create_default_protocol(): + protocol = HybridTopProtocol(settings=HybridTopProtocol.default_settings()) + + assert protocol + + +def test_serialize_protocol(): + protocol = HybridTopProtocol( + settings=HybridTopProtocol.default_settings(), + ) + + ser = protocol.to_dict() + + ret = HybridTopProtocol.from_dict(ser) + + assert protocol == ret + + +def test_position_overlap_fail(): + mapping = {"old_to_new_env_atom_map": {i: i for i in range(4)}} + positionsA = np.array([[1, 1, 1]] * 4) * unit.angstrom + positionsB = np.array([[1, 1, 2.1]] * 4) * unit.angstrom + + with pytest.raises(ValueError, match="deviates by more than"): + HybridTopProtocolUnit._check_position_overlap( + mapping, + positionsA, + positionsB, + ) + + +def test_position_overlap_warn(): + mapping = {"old_to_new_env_atom_map": {}, "old_to_new_core_atom_map": {i: i for i in range(4)}} + positionsA = np.array([[1, 1, 1]] * 4) * unit.angstrom + positionsB = np.array([[1, 1, 2.1]] * 4) * unit.angstrom + + with pytest.warns(UserWarning, match="deviates by more than"): + HybridTopProtocolUnit._check_position_overlap( + mapping, + positionsA, + positionsB, + ) + + +@pytest.mark.parametrize("method", ["repex", "sams", "independent", "InDePeNdENT"]) +def test_dry_run_default_vacuum( + benzene_vacuum_system, toluene_vacuum_system, benzene_to_toluene_mapping, method, tmpdir +): + vac_settings = HybridTopProtocol.default_settings() + vac_settings.forcefield_settings.nonbonded_method = "nocutoff" + vac_settings.simulation_settings.sampler_method = method + vac_settings.protocol_repeats = 1 + + protocol = HybridTopProtocol( + settings=vac_settings, + ) + + # create DAG from protocol and take first (and only) work unit from within + dag = protocol.create( + stateA=benzene_vacuum_system, + stateB=toluene_vacuum_system, + mapping=benzene_to_toluene_mapping, + ) + dag_unit = list(dag.protocol_units)[0] + + with tmpdir.as_cwd(): + sampler = dag_unit.run(dry=True)["debug"]["sampler"] + assert isinstance(sampler, MultiStateSampler) + assert not sampler.is_periodic + assert sampler._thermodynamic_states[0].barostat is None + + # Check hybrid OMM and MDTtraj Topologies + htf = sampler._hybrid_factory + # 16 atoms: + # 11 common atoms, 1 extra hydrogen in benzene, 4 extra in toluene + # 12 bonds in benzene + 4 extra toluene bonds + assert len(list(htf.hybrid_topology.atoms)) == 16 + assert len(list(htf.omm_hybrid_topology.atoms())) == 16 + assert len(list(htf.hybrid_topology.bonds)) == 16 + assert len(list(htf.omm_hybrid_topology.bonds())) == 16 + + # smoke test - can convert back the mdtraj topology + ret_top = mdt.Topology.to_openmm(htf.hybrid_topology) + assert len(list(ret_top.atoms())) == 16 + assert len(list(ret_top.bonds())) == 16 + + # check that our PDB has the right number of atoms + pdb = mdt.load_pdb("hybrid_system.pdb") + assert pdb.n_atoms == 16 + + # check the system forces + system = htf.hybrid_system + assert len(system.getForces()) == 9 + + def assert_force_num(system, forcetype, number): + forces = [f for f in system.getForces() if isinstance(f, forcetype)] + assert len(forces) == number + + assert_force_num(system, NonbondedForce, 1) + assert_force_num(system, CustomNonbondedForce, 1) + assert_force_num(system, CustomBondForce, 2) + assert_force_num(system, CustomAngleForce, 1) + assert_force_num(system, CustomTorsionForce, 1) + assert_force_num(system, HarmonicBondForce, 1) + assert_force_num(system, HarmonicAngleForce, 1) + assert_force_num(system, PeriodicTorsionForce, 1) + + # Check the nonbonded force is NoCutoff + nonbond = [f for f in system.getForces() if isinstance(f, NonbondedForce)] + assert nonbond[0].getNonbondedMethod() == NonbondedForce.NoCutoff + + +BENZ = """\ +benzene + PyMOL2.5 3D 0 + + 12 12 0 0 0 0 0 0 0 0999 V2000 + 1.4045 -0.0000 0.0000 C 0 0 0 0 0 0 0 0 0 0 0 0 + 0.7022 1.2164 0.0000 C 0 0 0 0 0 0 0 0 0 0 0 0 + -0.7023 1.2164 0.0000 C 0 0 0 0 0 0 0 0 0 0 0 0 + -1.4045 -0.0000 0.0000 C 0 0 0 0 0 0 0 0 0 0 0 0 + -0.7023 -1.2164 0.0000 C 0 0 0 0 0 0 0 0 0 0 0 0 + 0.7023 -1.2164 0.0000 C 0 0 0 0 0 0 0 0 0 0 0 0 + 2.5079 -0.0000 0.0000 H 0 0 0 0 0 0 0 0 0 0 0 0 + 1.2540 2.1720 0.0000 H 0 0 0 0 0 0 0 0 0 0 0 0 + -1.2540 2.1720 0.0000 H 0 0 0 0 0 0 0 0 0 0 0 0 + -2.5079 -0.0000 0.0000 H 0 0 0 0 0 0 0 0 0 0 0 0 + -1.2540 -2.1719 0.0000 H 0 0 0 0 0 0 0 0 0 0 0 0 + 1.2540 -2.1720 0.0000 H 0 0 0 0 0 0 0 0 0 0 0 0 + 1 2 2 0 0 0 0 + 1 6 1 0 0 0 0 + 1 7 1 0 0 0 0 + 2 3 1 0 0 0 0 + 2 8 1 0 0 0 0 + 3 4 2 0 0 0 0 + 3 9 1 0 0 0 0 + 4 5 1 0 0 0 0 + 4 10 1 0 0 0 0 + 5 6 2 0 0 0 0 + 5 11 1 0 0 0 0 + 6 12 1 0 0 0 0 +M END +$$$$ +""" + + +PYRIDINE = """\ +pyridine + PyMOL2.5 3D 0 + + 11 11 0 0 0 0 0 0 0 0999 V2000 + 1.4045 -0.0000 0.0000 C 0 0 0 0 0 0 0 0 0 0 0 0 + -0.7023 1.2164 0.0000 C 0 0 0 0 0 0 0 0 0 0 0 0 + -1.4045 -0.0000 0.0000 C 0 0 0 0 0 0 0 0 0 0 0 0 + -0.7023 -1.2164 0.0000 C 0 0 0 0 0 0 0 0 0 0 0 0 + 0.7023 -1.2164 0.0000 C 0 0 0 0 0 0 0 0 0 0 0 0 + 2.4940 -0.0325 0.0000 H 0 0 0 0 0 0 0 0 0 0 0 0 + 1.2473 -2.1604 0.0000 H 0 0 0 0 0 0 0 0 0 0 0 0 + -1.2473 -2.1604 0.0000 H 0 0 0 0 0 0 0 0 0 0 0 0 + -2.4945 -0.0000 0.0000 H 0 0 0 0 0 0 0 0 0 0 0 0 + -1.2753 2.1437 0.0000 H 0 0 0 0 0 0 0 0 0 0 0 0 + 0.7525 1.3034 0.0000 N 0 0 0 0 0 0 0 0 0 0 0 0 + 1 5 1 0 0 0 0 + 1 6 1 0 0 0 0 + 1 11 2 0 0 0 0 + 2 3 2 0 0 0 0 + 2 10 1 0 0 0 0 + 3 4 1 0 0 0 0 + 3 9 1 0 0 0 0 + 4 5 2 0 0 0 0 + 4 8 1 0 0 0 0 + 5 7 1 0 0 0 0 + 2 11 1 0 0 0 0 +M END +$$$$ +""" + + +def test_dry_core_element_change(tmpdir): + benz = openfe.SmallMoleculeComponent(Chem.MolFromMolBlock(BENZ, removeHs=False)) + pyr = openfe.SmallMoleculeComponent(Chem.MolFromMolBlock(PYRIDINE, removeHs=False)) + + mapping = openfe.LigandAtomMapping( + benz, pyr, {0: 0, 1: 10, 2: 1, 3: 2, 4: 3, 5: 4, 6: 5, 8: 9, 9: 8, 10: 7, 11: 6} + ) + + settings = HybridTopProtocol.default_settings() + settings.forcefield_settings.nonbonded_method = "nocutoff" + + protocol = HybridTopProtocol( + settings=settings, + ) + + dag = protocol.create( + stateA=openfe.ChemicalSystem( + { + "ligand": benz, + } + ), + stateB=openfe.ChemicalSystem( + { + "ligand": pyr, + } + ), + mapping=mapping, + ) + + dag_unit = list(dag.protocol_units)[0] + + with tmpdir.as_cwd(): + sampler = dag_unit.run(dry=True)["debug"]["sampler"] + system = sampler._hybrid_factory.hybrid_system + assert system.getNumParticles() == 12 + # Average mass between nitrogen and carbon + assert pytest.approx(system.getParticleMass(1)._value) == 12.0008030 + + # Get out the CustomNonbondedForce + cnf = [f for f in system.getForces() if f.__class__.__name__ == "CustomNonbondedForce"][0] + # there should be no new unique atoms + assert cnf.getInteractionGroupParameters(6) == [(), ()] + # there should be one old unique atom (spare hydrogen from the benzene) + assert cnf.getInteractionGroupParameters(7) == [(7,), (7,)] + + +def test_dry_run_ligand(benzene_system, toluene_system, benzene_to_toluene_mapping, tmpdir): + # this might be a bit time consuming + settings = HybridTopProtocol.default_settings() + settings.protocol_repeats = 1 + settings.output_settings.output_indices = "resname AAA" + + protocol = HybridTopProtocol( + settings=settings, + ) + dag = protocol.create( + stateA=benzene_system, + stateB=toluene_system, + mapping=benzene_to_toluene_mapping, + ) + dag_unit = list(dag.protocol_units)[0] + + with tmpdir.as_cwd(): + sampler = dag_unit.run(dry=True)["debug"]["sampler"] + assert isinstance(sampler, MultiStateSampler) + assert sampler.is_periodic + assert isinstance(sampler._thermodynamic_states[0].barostat, MonteCarloBarostat) + assert sampler._thermodynamic_states[1].pressure == 1 * omm_unit.bar + + # Check we have the right number of atoms in the PDB + pdb = mdt.load_pdb("hybrid_system.pdb") + assert pdb.n_atoms == 16 + + # Check system forces + system = sampler._hybrid_factory.hybrid_system + assert len(system.getForces()) == 10 + + def assert_force_num(system, forcetype, number): + forces = [f for f in system.getForces() if isinstance(f, forcetype)] + assert len(forces) == number + + assert_force_num(system, NonbondedForce, 1) + assert_force_num(system, CustomNonbondedForce, 1) + assert_force_num(system, CustomBondForce, 2) + assert_force_num(system, CustomAngleForce, 1) + assert_force_num(system, CustomTorsionForce, 1) + assert_force_num(system, HarmonicBondForce, 1) + assert_force_num(system, HarmonicAngleForce, 1) + assert_force_num(system, PeriodicTorsionForce, 1) + assert_force_num(system, MonteCarloBarostat, 1) + + # Check the nonbonded force is NoCutoff + nonbond = [f for f in system.getForces() if isinstance(f, NonbondedForce)] + assert nonbond[0].getNonbondedMethod() == NonbondedForce.PME + + +def test_dry_run_vacuum_user_charges(benzene_modifications, tmpdir): + """ + Create a hybrid system with a set of fictitious user supplied charges + and ensure that they are properly passed through to the constructed + hybrid topology. + """ + vac_settings = HybridTopProtocol.default_settings() + vac_settings.forcefield_settings.nonbonded_method = "nocutoff" + vac_settings.protocol_repeats = 1 + + protocol = HybridTopProtocol( + settings=vac_settings, + ) + + def assign_fictitious_charges(offmol): + """ + Get a random array of fake partial charges (ints because why not) + that sums up to 0. Note that OpenFF will complain if you try to + create a molecule that has a total charge that is different from + the expected formal charge, hence we enforce a zero charge here. + """ + rand_arr = np.random.randint(1, 10, size=offmol.n_atoms) / 100 + rand_arr[-1] = -sum(rand_arr[:-1]) + return rand_arr * unit.elementary_charge + + def check_propchgs(smc, charge_array): + """ + Check that the partial charges we assigned to our offmol from which + the smc was constructed are present and the right ones. + """ + prop_chgs = smc.to_dict()["molprops"]["atom.dprop.PartialCharge"] + prop_chgs = np.array(prop_chgs.split(), dtype=float) + np.testing.assert_allclose(prop_chgs, charge_array.m) + + # Create new smc with overriden charges + benzene_offmol = benzene_modifications["benzene"].to_openff() + toluene_offmol = benzene_modifications["toluene"].to_openff() + benzene_rand_chg = assign_fictitious_charges(benzene_offmol) + toluene_rand_chg = assign_fictitious_charges(toluene_offmol) + benzene_offmol.partial_charges = benzene_rand_chg + toluene_offmol.partial_charges = toluene_rand_chg + benzene_smc = openfe.SmallMoleculeComponent.from_openff(benzene_offmol) + toluene_smc = openfe.SmallMoleculeComponent.from_openff(toluene_offmol) + + # Check that the new smcs have the new overriden charges + check_propchgs(benzene_smc, benzene_rand_chg) + check_propchgs(toluene_smc, toluene_rand_chg) + + # Create new mapping + mapper = openfe.setup.LomapAtomMapper(element_change=False) + mapping = next(mapper.suggest_mappings(benzene_smc, toluene_smc)) + + # create DAG from protocol and take first (and only) work unit from within + dag = protocol.create( + stateA=openfe.ChemicalSystem( + { + "l": benzene_smc, + } + ), + stateB=openfe.ChemicalSystem( + { + "l": toluene_smc, + } + ), + mapping=mapping, + ) + dag_unit = list(dag.protocol_units)[0] + + with tmpdir.as_cwd(): + sampler = dag_unit.run(dry=True)["debug"]["sampler"] + htf = sampler._factory + hybrid_system = htf.hybrid_system + + # get the standard nonbonded force + nonbond = [f for f in hybrid_system.getForces() if isinstance(f, NonbondedForce)] + assert len(nonbond) == 1 + + # get the particle parameter offsets + c_offsets = {} + for i in range(nonbond[0].getNumParticleParameterOffsets()): + offset = nonbond[0].getParticleParameterOffset(i) + c_offsets[offset[1]] = ensure_quantity(offset[2], "openff") + + # Here is a bit of exposition on what we're doing + # HTF creates two sets of nonbonded forces, a standard one (for the + # PME) and a custom one (for sterics). + # Here we specifically check charges, so we only concentrate on the + # standard NonbondedForce. + # The way the NonbondedForce is constructed is as follows: + # - unique old atoms: + # * The particle charge is set to the input molA particle charge + # * The chargeScale offset is set to the negative value of the molA + # particle charge (such that by scaling you effectively zero out + # the charge. + # - unique new atoms: + # * The particle charge is set to zero (doesn't exist in the starting + # end state). + # * The chargeScale offset is set to the value of the molB particle + # charge (such that by scaling you effectively go from 0 to molB + # charge). + # - core atoms: + # * The particle charge is set to the input molA particle charge + # (i.e. we start from a system that has molA charges). + # * The particle charge offset is set to the difference between + # the molB particle charge and the molA particle charge (i.e. + # we scale by that difference to get to the value of the molB + # particle charge). + for i in range(hybrid_system.getNumParticles()): + c, s, e = nonbond[0].getParticleParameters(i) + # get the particle charge (c) + c = ensure_quantity(c, "openff") + # particle charge (c) is equal to molA particle charge + # offset (c_offsets) is equal to -(molA particle charge) + if i in htf._atom_classes["unique_old_atoms"]: + idx = htf._hybrid_to_old_map[i] + np.testing.assert_allclose(c, benzene_rand_chg[idx]) + np.testing.assert_allclose(c_offsets[i], -benzene_rand_chg[idx]) + # particle charge (c) is equal to 0 + # offset (c_offsets) is equal to molB particle charge + elif i in htf._atom_classes["unique_new_atoms"]: + idx = htf._hybrid_to_new_map[i] + np.testing.assert_allclose(c, 0 * unit.elementary_charge) + np.testing.assert_allclose(c_offsets[i], toluene_rand_chg[idx]) + # particle charge (c) is equal to molA particle charge + # offset (c_offsets) is equal to difference between molB and molA + elif i in htf._atom_classes["core_atoms"]: + old_i = htf._hybrid_to_old_map[i] + new_i = htf._hybrid_to_new_map[i] + c_exp = toluene_rand_chg[new_i] - benzene_rand_chg[old_i] + np.testing.assert_allclose(c, benzene_rand_chg[old_i]) + np.testing.assert_allclose(c_offsets[i], c_exp) diff --git a/src/pontibus/tests/protocols/relative/test_protocol_slow.py b/src/pontibus/tests/protocols/relative/test_protocol_slow.py new file mode 100644 index 0000000..16954a9 --- /dev/null +++ b/src/pontibus/tests/protocols/relative/test_protocol_slow.py @@ -0,0 +1,82 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe + +import pathlib + +import pytest +from gufe.protocols import execute_DAG +from openff.units import unit + +from pontibus.protocols.relative import HybridTopProtocol + + +@pytest.mark.gpu +def test_vacuum(benzene_vacuum_system, toluene_vacuum_system, benzene_to_toluene_mapping, tmpdir): + """ + Run a short MD simulation and make sure things didn't fail. + """ + s = HybridTopProtocol.default_settings() + s.simulation_settings.equilibration_length = 100 * unit.picosecond + s.simulation_settings.production_length = 500 * unit.picosecond + s.forcefield_settings.nonbonded_method = "nocutoff" + s.protocol_repeats = 1 + s.engine_settings.compute_platform = "CUDA" + + p = HybridTopProtocol(s) + + dag = p.create( + stateA=benzene_vacuum_system, + stateB=toluene_vacuum_system, + mapping=benzene_to_toluene_mapping, + ) + + cwd = pathlib.Path(str(tmpdir)) + r = execute_DAG(dag, shared_basedir=cwd, scratch_basedir=cwd, keep_shared=True) + + assert r.ok() + for pur in r.protocol_unit_results: + unit_shared = tmpdir / f"shared_{pur.source_key}_attempt_0" + assert unit_shared.exists() + assert pathlib.Path(unit_shared).is_dir() + + # Check the checkpoint file exists + checkpoint = pur.outputs["last_checkpoint"] + assert checkpoint == "checkpoint.chk" + assert (unit_shared / checkpoint).exists() + + # Check the nc simulation file exists + nc = pur.outputs["nc"] + assert nc == unit_shared / "simulation.nc" + assert nc.exists() + + # Check structural analysis contents + # TODO: for now this is disabled due to issue #117 + # structural_analysis_file = unit_shared / "structural_analysis.npz" + # assert (structural_analysis_file).exists() + # assert pur.outputs['structural_analysis'] == structural_analysis_file + + # structural_data = np.load(pur.outputs['structural_analysis']) + # structural_keys = [ + # 'protein_RMSD', 'ligand_RMSD', 'ligand_COM_drift', + # 'protein_2D_RMSD', 'time_ps' + # ] + # for key in structural_keys: + # assert key in structural_data.keys() + + # 6 frames being written to file + # assert_allclose(structural_data['time_ps'], [0.0, 0.02, 0.04, 0.06, 0.08, 0.1]) + # assert structural_data['ligand_RMSD'].shape == (11, 6) + # assert structural_data['ligand_COM_drift'].shape == (11, 6) + # No protein so should be empty + # assert structural_data['protein_RMSD'].size == 0 + # assert structural_data['protein_2D_RMSD'].size == 0 + + # Test results + results = p.gather([r]) + estimate = results.get_estimate() + assert estimate.m == pytest.approx(0.80, abs=0.2) + uncert = results.get_uncertainty() + assert uncert.m == pytest.approx(0.0) + states = results.get_replica_states() + assert len(states) == 1 + assert states[0].shape[1] == 11 diff --git a/src/pontibus/tests/protocols/relative/test_results.py b/src/pontibus/tests/protocols/relative/test_results.py new file mode 100644 index 0000000..7e7157b --- /dev/null +++ b/src/pontibus/tests/protocols/relative/test_results.py @@ -0,0 +1,173 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe +import json + +import gufe +import numpy as np +import openfe +import pytest +from openff.units import unit as offunit + +from pontibus.protocols.relative import HybridTopProtocolResult + + +class TestSolventProtocolResult: + @pytest.fixture() + def protocolresult(self, rfe_solv_transformation_json): + d = json.loads( + rfe_solv_transformation_json, + cls=gufe.tokenization.JSON_HANDLER.decoder, + ) + + pr = openfe.ProtocolResult.from_dict(d["protocol_result"]) + + return pr + + def test_reload_protocol_result(self, rfe_solv_transformation_json): + d = json.loads( + rfe_solv_transformation_json, + cls=gufe.tokenization.JSON_HANDLER.decoder, + ) + + pr = HybridTopProtocolResult.from_dict(d["protocol_result"]) + + assert pr + + def test_get_estimate(self, protocolresult): + est = protocolresult.get_estimate() + + assert est + assert est.m == pytest.approx(16.94, abs=0.5) + assert isinstance(est, offunit.Quantity) + assert est.is_compatible_with(offunit.kilojoule_per_mole) + + def test_get_uncertainty(self, protocolresult): + est = protocolresult.get_uncertainty() + + assert est + assert est.m == pytest.approx(0.2, abs=0.2) + assert isinstance(est, offunit.Quantity) + assert est.is_compatible_with(offunit.kilojoule_per_mole) + + def test_get_individual(self, protocolresult): + inds = protocolresult.get_individual_estimates() + + assert isinstance(inds, list) + assert len(inds) == 3 + + for e, u in inds: + assert e.is_compatible_with(offunit.kilojoule_per_mole) + assert u.is_compatible_with(offunit.kilojoule_per_mole) + + def test_get_forwards_etc(self, protocolresult): + far = protocolresult.get_forward_and_reverse_energy_analysis() + + assert isinstance(far, list) + far1 = far[0] + assert isinstance(far1, dict) + + for k in [ + "fractions", + "forward_DGs", + "forward_dDGs", + "reverse_DGs", + "reverse_dDGs", + ]: + assert k in far1 + + if k == "fractions": + assert isinstance(far1[k], np.ndarray) + else: + assert isinstance(far1[k], offunit.Quantity) + assert far1[k].is_compatible_with(offunit.kilojoule_per_mole) + + def test_get_frwd_reverse_none_return(self, protocolresult): + # fetch the first result + data = [i for i in protocolresult.data.values()][0][0] + # set the output to None + data.outputs["forward_and_reverse_energies"] = None + + # now fetch the analysis results and expect a warning + wmsg = "One or more ``None`` entries were found in" + with pytest.warns(UserWarning, match=wmsg): + protocolresult.get_forward_and_reverse_energy_analysis() + + def test_get_overlap_matrices(self, protocolresult): + ovp = protocolresult.get_overlap_matrices() + + assert isinstance(ovp, list) + assert len(ovp) == 3 + + ovp1 = ovp[0] + assert isinstance(ovp1["matrix"], np.ndarray) + assert ovp1["matrix"].shape == (11, 11) + + def test_get_replica_transition_statistics(self, protocolresult): + rpx = protocolresult.get_replica_transition_statistics() + + assert isinstance(rpx, list) + assert len(rpx) == 3 + rpx1 = rpx[0] + assert "eigenvalues" in rpx1 + assert "matrix" in rpx1 + assert rpx1["eigenvalues"].shape == (11,) + assert rpx1["matrix"].shape == (11, 11) + + def test_equilibration_iterations(self, protocolresult): + eq = protocolresult.equilibration_iterations() + + assert isinstance(eq, list) + assert len(eq) == 3 + assert all(isinstance(v, float) for v in eq) + + def test_production_iterations(self, protocolresult): + prod = protocolresult.production_iterations() + + assert isinstance(prod, list) + assert len(prod) == 3 + assert all(isinstance(v, float) for v in prod) + + def test_filenotfound_replica_states(self, protocolresult): + errmsg = "File could not be found" + + with pytest.raises(ValueError, match=errmsg): + protocolresult.get_replica_states() + + +class TestVacuumProtocolResult(TestSolventProtocolResult): + @pytest.fixture() + def protocolresult(self, rfe_vacuum_transformation_json): + d = json.loads( + rfe_vacuum_transformation_json, + cls=gufe.tokenization.JSON_HANDLER.decoder, + ) + + pr = openfe.ProtocolResult.from_dict(d["protocol_result"]) + + return pr + + def test_reload_protocol_result(self, rfe_vacuum_transformation_json): + d = json.loads( + rfe_vacuum_transformation_json, + cls=gufe.tokenization.JSON_HANDLER.decoder, + ) + + pr = HybridTopProtocolResult.from_dict(d["protocol_result"]) + + assert pr + + def test_get_estimate(self, protocolresult): + est = protocolresult.get_estimate() + + assert est + assert est.m == pytest.approx(16.94, abs=0.5) + assert isinstance(est, offunit.Quantity) + assert est.is_compatible_with(offunit.kilojoule_per_mole) + + def test_get_uncertainty(self, protocolresult): + est = protocolresult.get_uncertainty() + + assert est + assert est.m == pytest.approx(0.16, abs=0.2) + assert isinstance(est, offunit.Quantity) + assert est.is_compatible_with(offunit.kilojoule_per_mole) diff --git a/src/pontibus/tests/protocols/relative/test_tokenization.py b/src/pontibus/tests/protocols/relative/test_tokenization.py new file mode 100644 index 0000000..4137eef --- /dev/null +++ b/src/pontibus/tests/protocols/relative/test_tokenization.py @@ -0,0 +1,128 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe +import json + +import gufe +import pytest +from gufe.tests.test_tokenization import GufeTokenizableTestsMixin + +from pontibus.protocols.relative import ( + HybridTopProtocol, + HybridTopProtocolResult, + HybridTopProtocolUnit, +) + + +@pytest.fixture +def protocol(): + return HybridTopProtocol(HybridTopProtocol.default_settings()) + + +@pytest.fixture +def vacuum_protocol(): + settings = HybridTopProtocol.default_settings() + settings.forcefield_settings.nonbonded_method = "nocutoff" + return HybridTopProtocol(settings=settings) + + +@pytest.fixture +def vacuum_protocol_unit( + vacuum_protocol, benzene_to_toluene_mapping, benzene_vacuum_system, toluene_vacuum_system +): + pus = vacuum_protocol.create( + stateA=benzene_vacuum_system, + stateB=toluene_vacuum_system, + mapping=benzene_to_toluene_mapping, + ) + return list(pus.protocol_units)[0] + + +@pytest.fixture +def solvent_protocol_unit(protocol, benzene_to_toluene_mapping, benzene_system, toluene_system): + pus = protocol.create( + stateA=benzene_system, + stateB=toluene_system, + mapping=benzene_to_toluene_mapping, + ) + return list(pus.protocol_units)[0] + + +@pytest.fixture +def protocol_result(rfe_solv_transformation_json): + d = json.loads(rfe_solv_transformation_json, cls=gufe.tokenization.JSON_HANDLER.decoder) + pr = HybridTopProtocolResult.from_dict(d["protocol_result"]) + return pr + + +class TestProtocol(GufeTokenizableTestsMixin): + cls = HybridTopProtocol + key = None + repr = "HybridTopProtocol-" + + @pytest.fixture() + def instance(self, protocol): + return protocol + + def test_repr(self, instance): + """ + Overwrites the base `test_repr` call. + """ + assert isinstance(repr(instance), str) + assert self.repr in repr(instance) + + +class TestSolventUnit(GufeTokenizableTestsMixin): + cls = HybridTopProtocolUnit + repr = "HybridTopProtocolUnit(" + key = None + + @pytest.fixture() + def instance(self, solvent_protocol_unit): + return solvent_protocol_unit + + def test_key_stable(self): + pytest.skip() + + def test_repr(self, instance): + """ + Overwrites the base `test_repr` call. + """ + assert isinstance(repr(instance), str) + assert self.repr in repr(instance) + + +class TestVacuumUnit(GufeTokenizableTestsMixin): + cls = HybridTopProtocolUnit + repr = "HybridTopProtocolUnit(" + key = None + + @pytest.fixture() + def instance(self, vacuum_protocol_unit): + return vacuum_protocol_unit + + def test_key_stable(self): + pytest.skip() + + def test_repr(self, instance): + """ + Overwrites the base `test_repr` call. + """ + assert isinstance(repr(instance), str) + assert self.repr in repr(instance) + + +class TestProtocolResult(GufeTokenizableTestsMixin): + cls = HybridTopProtocolResult + key = None + repr = "HybridTopProtocolResult-" + + @pytest.fixture() + def instance(self, protocol_result): + return protocol_result + + def test_repr(self, instance): + """ + Overwrites the base `test_repr` call. + """ + assert isinstance(repr(instance), str) + assert self.repr in repr(instance) diff --git a/src/pontibus/tests/protocols/solvation/conftest.py b/src/pontibus/tests/protocols/solvation/conftest.py index fc911b8..89f6069 100644 --- a/src/pontibus/tests/protocols/solvation/conftest.py +++ b/src/pontibus/tests/protocols/solvation/conftest.py @@ -17,7 +17,7 @@ def afe_solv_water_transformation_json() -> str: """ ASFE results object as created by quickrun. - generated with devtools/gent-serialized-results.py + generated with devtools/gen-serialized-results.py """ d = resources.files("pontibus.tests.data.solvation_protocol") file = d / "ASFEProtocol_water_json_results.gz" @@ -31,7 +31,7 @@ def afe_solv_octanol_transformation_json() -> str: """ ASFE results object as created by quickrun. - generated with devtools/gent-serialized-results.py + generated with devtools/gen-serialized-results.py """ d = resources.files("pontibus.tests.data.solvation_protocol") file = d / "ASFEProtocol_octanol_json_results.gz" diff --git a/src/pontibus/tests/protocols/solvation/test_dry_run.py b/src/pontibus/tests/protocols/solvation/test_dry_run.py index b018bc0..9321ee1 100644 --- a/src/pontibus/tests/protocols/solvation/test_dry_run.py +++ b/src/pontibus/tests/protocols/solvation/test_dry_run.py @@ -66,7 +66,6 @@ def test_dry_run_vacuum_benzene(charged_benzene, method, tmpdir): assert not vac_sampler.is_periodic system = vac_sampler._thermodynamic_states[0].get_system(remove_thermostat=True) - print(system.getForces()) assert len(system.getForces()) == 12 def assert_force_num(system, forcetype, number): diff --git a/src/pontibus/tests/utils/test_interchange_packmol.py b/src/pontibus/tests/utils/test_interchange_packmol.py index 96dcd6b..247e5a2 100644 --- a/src/pontibus/tests/utils/test_interchange_packmol.py +++ b/src/pontibus/tests/utils/test_interchange_packmol.py @@ -23,14 +23,16 @@ InterchangeFFSettings, PackmolSolvationSettings, ) +from pontibus.utils.molecule_utils import ( + _check_library_charges, + _get_offmol_resname, + _set_offmol_resname, +) from pontibus.utils.molecules import WATER from pontibus.utils.system_creation import ( _check_and_deduplicate_charged_mols, - _check_library_charges, _get_comp_resnames, _get_force_field, - _get_offmol_resname, - _set_offmol_resname, _solvate_system, interchange_packmol_creation, ) @@ -105,7 +107,7 @@ def test_get_and_set_offmol_resname(CN_molecule, caplog): with caplog.at_level(logging.WARNING): assert _get_offmol_resname(CN_off) is None - assert "Inconsistent residue name" in caplog.text + assert "Inconsistent metadata residue_name" in caplog.text def test_check_library_charges_pass(water_off): diff --git a/src/pontibus/tests/utils/test_system_manipulation.py b/src/pontibus/tests/utils/test_system_manipulation.py new file mode 100644 index 0000000..5e52beb --- /dev/null +++ b/src/pontibus/tests/utils/test_system_manipulation.py @@ -0,0 +1,145 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe + +import pytest +from numpy.testing import assert_allclose +from openff.interchange import Interchange +from openff.interchange.components._packmol import solvate_topology +from openff.toolkit import ForceField, Molecule, Topology +from openmm import CMMotionRemover, MonteCarloBarostat, System +from openmm import unit as omm_unit + +from pontibus.utils.molecule_utils import ( + _get_offmol_metadata, + _set_offmol_metadata, +) +from pontibus.utils.system_manipulation import ( + adjust_system, + copy_interchange_with_replacement, +) + + +def test_adjust_forces_nothing(): + """ + A smoke test, this should just pass. + """ + system = System() + adjust_system(system) + + +def test_ajdust_forces_remove_com_remover(): + system = System() + com_force = CMMotionRemover() + system.addForce(com_force) + adjust_system(system, remove_force_types=CMMotionRemover) + + assert system.getNumForces() == 0 + + +def test_adjust_forces_add_comm_and_barostat(): + system = System() + barostat = MonteCarloBarostat(1.0 * omm_unit.bar, 298.15 * omm_unit.kelvin) + com_force = CMMotionRemover() + adjust_system(system, add_forces=[barostat, com_force]) + + assert system.getNumForces() == 2 + + +@pytest.fixture(scope="module") +def forcefield(): + return ForceField("openff-2.2.1.offxml") + + +@pytest.fixture(scope="module") +def insert_molecule(): + m = Molecule.from_smiles("CCCC") + m.generate_conformers(n_conformers=1) + return m + + +@pytest.fixture(scope="module") +def del_molecule(): + m = Molecule.from_smiles("CCO") + m.generate_conformers(n_conformers=1) + return m + + +def test_copy_no_conformers(forcefield): + m1 = Molecule.from_smiles("C") + m2 = Molecule.from_smiles("O") + topology = Topology.from_molecules([m1]) + inter = Interchange.from_smirnoff(forcefield, topology) + + with pytest.raises(ValueError, match="molecules need conformers"): + _ = copy_interchange_with_replacement( + interchange=inter, + del_mol=m1, + insert_mol=m2, + force_field=forcefield, + charged_molecules=None, + ) + + +def test_copy_equality_clash(forcefield, insert_molecule, del_molecule): + topology = Topology.from_molecules([del_molecule, del_molecule]) + inter = Interchange.from_smirnoff(forcefield, topology) + + with pytest.raises(ValueError, match="equality clash"): + _ = copy_interchange_with_replacement( + interchange=inter, + del_mol=del_molecule, + insert_mol=insert_molecule, + force_field=forcefield, + charged_molecules=None, + ) + + +def test_copy_no_del_match(forcefield, insert_molecule, del_molecule): + fake_del_mol = Molecule.from_smiles("C") + fake_del_mol.generate_conformers(n_conformers=1) + topology = Topology.from_molecules([del_molecule]) + inter = Interchange.from_smirnoff(forcefield, topology) + + with pytest.raises(ValueError, match="matching del_mol in input"): + _ = copy_interchange_with_replacement( + interchange=inter, + del_mol=fake_del_mol, + insert_mol=insert_molecule, + force_field=forcefield, + charged_molecules=None, + ) + + +def test_copy_full(forcefield): + m1 = Molecule.from_smiles("CCCC") + m1.generate_conformers(n_conformers=1) + m1.assign_partial_charges(partial_charge_method="gasteiger") + _set_offmol_metadata(m1, "residue_number", 999) + m2 = Molecule.from_smiles("CCCO") + m2.generate_conformers(n_conformers=1) + m2.assign_partial_charges(partial_charge_method="gasteiger") + + # Solvate m1 + solvated_top = solvate_topology(Topology.from_molecules([m1])) + + # Create interchange + inter = Interchange.from_smirnoff(forcefield, solvated_top, charge_from_molecules=[m1]) + + inter_new = copy_interchange_with_replacement( + interchange=inter, del_mol=m1, insert_mol=m2, force_field=forcefield, charged_molecules=[m2] + ) + + assert inter.topology.n_molecules == inter_new.topology.n_molecules + assert inter_new.topology.n_unique_molecules == 4 + for idx in range(inter_new.topology.n_molecules - 1): + mol_new = inter_new.topology.molecule(idx) + mol_old = inter.topology.molecule(idx + 1) + assert mol_new.is_isomorphic_with(mol_old) + + assert_allclose(mol_new.conformers[0], mol_old.conformers[0]) + + insert_mol_new = inter_new.topology.molecule(inter_new.topology.n_molecules - 1) + assert insert_mol_new.is_isomorphic_with(m2) + assert_allclose(m2.conformers[0], insert_mol_new.conformers[0]) + assert_allclose(m2.partial_charges, insert_mol_new.partial_charges) + assert _get_offmol_metadata(insert_mol_new, "residue_number") == 999 diff --git a/src/pontibus/utils/molecule_utils.py b/src/pontibus/utils/molecule_utils.py new file mode 100644 index 0000000..798edd8 --- /dev/null +++ b/src/pontibus/utils/molecule_utils.py @@ -0,0 +1,135 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe +import logging +from typing import Any + +from openff.toolkit import ForceField +from openff.toolkit import Molecule as OFFMolecule + +logger = logging.getLogger(__name__) + + +def _set_offmol_metadata( + offmol: OFFMolecule, + key: Any, + val: Any | None, +) -> None: + """ + Set a given metadata entry for a whole Molecule. + + Parameters + ---------- + offmol : openff.toolkit.Molecule + The Molecule to set the metadata for. + key : Any + The metadata key. + val : Any + The value to set the metadata entry to. + """ + if val is None: + for a in offmol.atoms: + a.metadata.pop(key, None) + else: + for a in offmol.atoms: + a.metadata[key] = val + + +def _get_offmol_metadata(offmol: OFFMolecule, key: Any) -> Any | None: + """ + Get an offmol's given metadata entry and make sure it is + consistent across all atoms in the Molecule. + + Parameters + ---------- + offmol : openff.toolkit.Molecule + Molecule to get the metadata value from. + key: Any + The metadata entry key. + + Returns + ------- + value : Any | None + Metadata for the given key in the molecule. ``None`` if the + Molecule does not have that metadata entry set, or if + the value is inconsistent across all the atoms. + """ + value: Any | None = None + for a in offmol.atoms: + if value is None: + try: + value = a.metadata[key] + except KeyError: + return None + + if value != a.metadata[key]: + wmsg = f"Inconsistent metadata {key} in OFFMol: {offmol}" + logger.warning(wmsg) + return None + + return value + + +def _set_offmol_resname( + offmol: OFFMolecule, + resname: str | None, +) -> None: + """ + Helper method to set offmol residue names + + Parameters + ---------- + offmol : openff.toolkit.Molecule + Molecule to assign a residue name to. + resname : str | None + Residue name to be set. Set to None to clear it. + + Returns + ------- + None + """ + _set_offmol_metadata(offmol, "residue_name", resname) + + +def _get_offmol_resname(offmol: OFFMolecule) -> str | None: + """ + Helper method to get an offmol's residue name and make sure it is + consistent across all atoms in the Molecule. + + Parameters + ---------- + offmol : openff.toolkit.Molecule + Molecule to get the residue name from. + + Returns + ------- + resname : Optional[str] + Residue name of the molecule. ``None`` if the Molecule + does not have a residue name, or if the residue name is + inconsistent across all the atoms. + """ + return _get_offmol_metadata(offmol, "residue_name") + + +def _check_library_charges( + force_field: ForceField, + offmol: OFFMolecule, +) -> None: + """ + Check that library charges exists for an input molecule. + + force_field : openff.toolkit.ForceField + Force Field object with library charges. + offmol : openff.toolkit.Molecule + Molecule to check for matching library charges. + + Raises + ------ + ValueError + If no library charges are found for the molecule. + """ + handler = force_field.get_parameter_handler("LibraryCharges") + matches = handler.find_matches(offmol.to_topology()) + + if len(matches) == 0: + errmsg = f"No library charges found for {offmol}" + raise ValueError(errmsg) diff --git a/src/pontibus/utils/protocol_utils.py b/src/pontibus/utils/protocol_utils.py new file mode 100644 index 0000000..fac1b0f --- /dev/null +++ b/src/pontibus/utils/protocol_utils.py @@ -0,0 +1,66 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe + +""" +Reusable methods for pontibus Protocols. +""" + +from openfe import SolventComponent +from openfe.protocols.openmm_utils import charge_generation +from openfe.protocols.openmm_utils.omm_settings import OpenFFPartialChargeSettings +from openff.toolkit import Molecule as OFFMolecule + +from pontibus.components.extended_solvent_component import ExtendedSolventComponent +from pontibus.utils.settings import PackmolSolvationSettings + + +def _get_and_charge_solvent_offmol( + solvent_component: SolventComponent | ExtendedSolventComponent, + solvation_settings: PackmolSolvationSettings, + partial_charge_settings: OpenFFPartialChargeSettings, +) -> OFFMolecule: + """ + Helper method to fetch the solvent offmol either + from an existing solvent_smcs, or from smiles. + + Parameters + ---------- + solvent_component : SolventComponent + smiles for the solvent molecule + solvation_settings : PackmolSolvationSettings + Settings defining how the system will be solvated + partial_charge_settings : OpenFFPartialChargeSettigns + Settings defining how partial charges are applied + + Returns + ------- + offmol : openff.toolkit.Molecule + + Notes + ----- + * If created from a smiles, the solvent will be assigned + a single conformer through `Molecule.generate_conformers`. + """ + # Get the solvent offmol + if isinstance(solvent_component, ExtendedSolventComponent): + solvent_offmol = solvent_component.solvent_molecule.to_openff() # type: ignore[union-attr] + else: + # If not, we create the solvent from smiles + # We generate a single conformer to avoid packing issues + solvent_offmol = OFFMolecule.from_smiles(solvent_component.smiles) + solvent_offmol.generate_conformers(n_conformers=1) + + # In-place assign solvent offmol charges if necessary + # Note: we don't enforce partial charge assignment to avoid + # cases where we want to rely on library charges instead. + if solvation_settings.assign_solvent_charges: + charge_generation.assign_offmol_partial_charges( + offmol=solvent_offmol, + overwrite=False, + method=partial_charge_settings.partial_charge_method, + toolkit_backend=partial_charge_settings.off_toolkit_backend, + generate_n_conformers=partial_charge_settings.number_of_conformers, + nagl_model=partial_charge_settings.nagl_model, + ) + + return solvent_offmol diff --git a/src/pontibus/utils/system_creation.py b/src/pontibus/utils/system_creation.py index 6ab3f66..f4fee22 100644 --- a/src/pontibus/utils/system_creation.py +++ b/src/pontibus/utils/system_creation.py @@ -20,95 +20,20 @@ from openff.toolkit import Molecule as OFFMolecule from openff.units import Quantity, unit -from pontibus.protocols.solvation.settings import ( +from pontibus.utils.molecule_utils import ( + _check_library_charges, + _get_offmol_resname, + _set_offmol_resname, +) +from pontibus.utils.molecules import offmol_water +from pontibus.utils.settings import ( InterchangeFFSettings, PackmolSolvationSettings, ) -from pontibus.utils.molecules import offmol_water logger = logging.getLogger(__name__) -def _set_offmol_resname( - offmol: OFFMolecule, - resname: str, -) -> None: - """ - Helper method to set offmol residue names - - Parameters - ---------- - offmol : openff.toolkit.Molecule - Molecule to assign a residue name to. - resname : str - Residue name to be set. - - Returns - ------- - None - """ - for a in offmol.atoms: - a.metadata["residue_name"] = resname - - -def _get_offmol_resname(offmol: OFFMolecule) -> str | None: - """ - Helper method to get an offmol's residue name and make sure it is - consistent across all atoms in the Molecule. - - Parameters - ---------- - offmol : openff.toolkit.Molecule - Molecule to get the residue name from. - - Returns - ------- - resname : Optional[str] - Residue name of the molecule. ``None`` if the Molecule - does not have a residue name, or if the residue name is - inconsistent across all the atoms. - """ - resname: str | None = None - for a in offmol.atoms: - if resname is None: - try: - resname = a.metadata["residue_name"] - except KeyError: - return None - - if resname != a.metadata["residue_name"]: - wmsg = f"Inconsistent residue name in OFFMol: {offmol} " - logger.warning(wmsg) - return None - - return resname - - -def _check_library_charges( - force_field: ForceField, - offmol: OFFMolecule, -) -> None: - """ - Check that library charges exists for an input molecule. - - force_field : openff.toolkit.ForceField - Force Field object with library charges. - offmol : openff.toolkit.Molecule - Molecule to check for matching library charges. - - Raises - ------ - ValueError - If no library charges are found for the molecule. - """ - handler = force_field.get_parameter_handler("LibraryCharges") - matches = handler.find_matches(offmol.to_topology()) - - if len(matches) == 0: - errmsg = f"No library charges found for {offmol}" - raise ValueError(errmsg) - - def _check_and_deduplicate_charged_mols( molecules: list[OFFMolecule], ) -> list[OFFMolecule]: @@ -438,7 +363,7 @@ def _solvate_system( continue if mol.is_isomorphic_with(solvent_offmol): - _set_offmol_resname(mol, _get_offmol_resname(solvent_offmol)) # type: ignore[arg-type] + _set_offmol_resname(mol, _get_offmol_resname(solvent_offmol)) if mol.is_isomorphic_with(na): _set_offmol_resname(mol, "NA+") diff --git a/src/pontibus/utils/system_manipulation.py b/src/pontibus/utils/system_manipulation.py new file mode 100644 index 0000000..2d631ed --- /dev/null +++ b/src/pontibus/utils/system_manipulation.py @@ -0,0 +1,130 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe + +import numpy as np +from openff.interchange import Interchange +from openff.toolkit import ForceField, Molecule, Topology +from openmm import Force, System + +from pontibus.utils.molecule_utils import ( + _get_offmol_metadata, + _set_offmol_metadata, +) +from pontibus.utils.system_creation import _check_and_deduplicate_charged_mols + + +def adjust_system( + system: System, + remove_force_types: type | list[type] | None = None, + add_forces: Force | list[Force] | None = None, +) -> None: + """ + Adjust a System by removing and adding forces as necessary. + + Parameters + ---------- + system : System + The OpenMM System to adjust + remove_force_types : type | list[type] | None + The force types to remove from the System, if present. + add_forces : list[Force] | None + The forces to add to the system. + """ + + def _adjust_inputs(var): + if var is not None: + if isinstance(var, list): + return var + return [var] + else: + return [] + + remove_force_types = _adjust_inputs(remove_force_types) + add_forces = _adjust_inputs(add_forces) + + for entry in remove_force_types: # type: ignore[union-attr] + for idx in reversed(range(system.getNumForces())): + force = system.getForce(idx) + if isinstance(force, entry): + system.removeForce(idx) + + for force in add_forces: + system.addForce(force) + + +def copy_interchange_with_replacement( + interchange: Interchange, + del_mol: Molecule, + insert_mol: Molecule, + force_field: ForceField, + charged_molecules: list[Molecule] | None, +) -> Interchange: + """ + Copy an Interchange deleting one Molecule and appending another. + + Parameters + ---------- + interchange : Interchange + Input Interchange to copy. + del_mol : Molecule + The Molecule to delete from the Interchange. + insert_mol : Molecule + The Molecule to insert to the Interchange. + force_field : ForceField + The ForceField object used to create the initial Interchange. + charged_molecules : list[Molecule] | None + A list of Molecules which partial charges to use in the new Interchange. + + Returns + ------- + new_interchange : Interchange + An copy of the input Interchange but with the Molecule mutation. + + Note + ---- + * ``del_mol`` is always deleted and ``insert_mol`` is appended to the end. + * The residue number of the Molecule matching ``del_mol`` in the input + Interchange is transcribed over to the ``insert_mol`` molecule. + """ + # Validate + if insert_mol.conformers is None or del_mol.conformers is None: + raise ValueError("Input molecules need conformers") + + # Get the del_mol idx + del_mol_idx = None + + # Search the Interchange Topology for a molecule with both + # isomorphic and spatial equality. + for idx, mol in enumerate(interchange.topology.molecules): + if mol.is_isomorphic_with(del_mol): + if np.allclose(mol.conformers[0], del_mol.conformers[0]): + if del_mol_idx is not None: + raise ValueError("equality clash with del_mol") + + del_mol_idx = idx + del_mol_resnum = _get_offmol_metadata(mol, "residue_number") + + if del_mol_idx is None: + errmsg = "No Molecule matching del_mol in input Interchange" + raise ValueError(errmsg) + + # Set molB residue number to molA + _set_offmol_metadata(insert_mol, "residue_number", del_mol_resnum) + + # Get a list of molecules from the input Interchange + mols = [m for m in interchange.topology.molecules] + mols.pop(del_mol_idx) # pop out the Molecule to be deleted + mols.append(insert_mol) # insert the new Molecule + + new_topology = Topology.from_molecules(mols) + new_topology.box_vectors = interchange.topology.box_vectors + + if charged_molecules is not None: + charged_molecules = _check_and_deduplicate_charged_mols(charged_molecules) + + new_interchange = force_field.create_interchange( + topology=new_topology, + charge_from_molecules=charged_molecules, + ) + + return new_interchange