diff --git a/src/openfe/protocols/openmm_afe/afe_protocol_results.py b/src/openfe/protocols/openmm_afe/afe_protocol_results.py index 10953f8fd..1563ae07d 100644 --- a/src/openfe/protocols/openmm_afe/afe_protocol_results.py +++ b/src/openfe/protocols/openmm_afe/afe_protocol_results.py @@ -31,8 +31,13 @@ class AbsoluteProtocolResultMixin: - bound_state = "solvent" - unbound_state = "vacuum" + """ + Subclasses must define the class attributes ``env_state`` and + ``ref_state``, naming the two legs of the thermodynamic cycle stored in + ``self.data`` (e.g. ``env_state = "complex"``, ``ref_state = "solvent"``). + """ + env_state: str + ref_state: str def __init__(self, **data): super().__init__(**data) @@ -40,7 +45,7 @@ def __init__(self, **data): if any( len(pur_list) > 2 for pur_list in itertools.chain( - self.data[self.bound_state].values(), self.data[self.unbound_state].values() + self.data[self.env_state].values(), self.data[self.ref_state].values() ) ): raise NotImplementedError("Can't stitch together results yet") @@ -49,7 +54,8 @@ def get_forward_and_reverse_energy_analysis( self, ) -> dict[str, list[Optional[dict[str, Union[npt.NDArray, Quantity]]]]]: """ - Get the reverse and forward analysis of the free energies. + Get the reverse and forward analysis of the free energies for both + legs of the thermodynamic cycle. Returns ------- @@ -57,7 +63,7 @@ def get_forward_and_reverse_energy_analysis( A dictionary, keyed for each leg of the thermodynamic cycle, either ``solvent`` and ``vacuum` for a solvation free energy or ``solvent`` and ``complex`` for a binding free energy, - with each containing a list of dictionaries containing the forward + with each containing a list of dictionaries with the forward and reverse analysis of each repeat of that simulation type. The forward and reverse analysis dictionaries contain: @@ -87,7 +93,7 @@ def get_forward_and_reverse_energy_analysis( forward_reverse: dict[str, list[Optional[dict[str, Union[npt.NDArray, Quantity]]]]] = {} - for key in [self.bound_state, self.unbound_state]: + for key in [self.env_state, self.ref_state]: forward_reverse[key] = [ pus[0].outputs["forward_and_reverse_energies"] for pus in self.data[key].values() # type: ignore[attr-defined] @@ -128,7 +134,7 @@ def get_overlap_matrices(self) -> dict[str, list[dict[str, npt.NDArray]]]: # Loop through and get the repeats and get the matrices overlap_stats: dict[str, list[dict[str, npt.NDArray]]] = {} - for key in [self.bound_state, self.unbound_state]: + for key in [self.env_state, self.ref_state]: overlap_stats[key] = [ pus[0].outputs["unit_mbar_overlap"] for pus in self.data[key].values() # type: ignore[attr-defined] @@ -138,8 +144,8 @@ def get_overlap_matrices(self) -> dict[str, list[dict[str, npt.NDArray]]]: def get_replica_transition_statistics(self) -> dict[str, list[dict[str, npt.NDArray]]]: """ - Get the replica exchange transition statistics for all - legs of the simulation. + Get the replica exchange transition statistics for both + legs of the thermodynamic cycle. Note ---- @@ -163,7 +169,7 @@ def get_replica_transition_statistics(self) -> dict[str, list[dict[str, npt.NDAr """ repex_stats: dict[str, list[dict[str, npt.NDArray]]] = {} try: - for key in [self.bound_state, self.unbound_state]: + for key in [self.env_state, self.ref_state]: repex_stats[key] = [ pus[0].outputs["replica_exchange_statistics"] for pus in self.data[key].values() # type: ignore[attr-defined] @@ -182,14 +188,14 @@ def get_replica_states(self) -> dict[str, list[npt.NDArray]]: ------- replica_states : dict[str, list[npt.NDArray]] Dictionary keyed for each leg of the thermodynamic cycle, either - `solvent` and `vacuum` for solvation free energies, - or `complex` and `solvent` for binding free energies, + ``solvent`` and ``vacuum`` for solvation free energies, + or ``complex`` and ``solvent`` for binding free energies, with lists of replica states timeseries for each repeat of that simulation type. """ replica_states: dict[str, list[npt.NDArray]] = { - self.bound_state: [], - self.unbound_state: [], + self.env_state: [], + self.ref_state: [], } def is_file(filename: str): @@ -215,7 +221,7 @@ def get_replica_state(nc, chk): return retval - for key in [self.bound_state, self.unbound_state]: + for key in [self.env_state, self.ref_state]: for pus in self.data[key].values(): # type: ignore[attr-defined] states = get_replica_state( pus[0].outputs["trajectory"], @@ -227,20 +233,21 @@ def get_replica_state(nc, chk): def equilibration_iterations(self) -> dict[str, list[float]]: """ - Get the number of equilibration iterations for each simulation. + Get the number of equilibration iterations for each repeat of + both legs of the calculation. Returns ------- equilibration_lengths : dict[str, list[float]] Dictionary keyed for each leg of the thermodynamic cycle, either - `solvent` and `vacuum` for solvation free energies, - or `complex` and `solvent` for binding free energies, + ``solvent`` and ``vacuum`` for solvation free energies, + or ``complex`` and ``solvent`` for binding free energies, with lists containing the number of equilibration iterations for each repeat of that simulation type. """ equilibration_lengths: dict[str, list[float]] = {} - for key in [self.bound_state, self.unbound_state]: + for key in [self.env_state, self.ref_state]: equilibration_lengths[key] = [ pus[0].outputs["equilibration_iterations"] for pus in self.data[key].values() # type: ignore[attr-defined] @@ -250,22 +257,21 @@ def equilibration_iterations(self) -> dict[str, list[float]]: def production_iterations(self) -> dict[str, list[float]]: """ - Get the number of production iterations for each simulation. Returns the number of uncorrelated production samples for each - repeat of the calculation. + repeat of both legs of the calculation. Returns ------- production_lengths : dict[str, list[float]] Dictionary keyed for each leg of the thermodynamic cycle, either - `solvent` and `vacuum` for solvation free energies, - or `complex` and `solvent` for binding free energies, + ``solvent`` and ``vacuum`` for solvation free energies, + or ``complex`` and ``solvent`` for binding free energies, with lists containing the number of equilibration iterations for each repeat of that simulation type. """ production_lengths: dict[str, list[float]] = {} - for key in [self.bound_state, self.unbound_state]: + for key in [self.env_state, self.ref_state]: production_lengths[key] = [ pus[0].outputs["production_iterations"] for pus in self.data[key].values() # type: ignore[attr-defined] @@ -282,15 +288,15 @@ def selection_indices(self) -> dict[str, list[Optional[npt.NDArray]]]: ------- indices : dict[str, list[npt.NDArray]] A dictionary keyed for each state, either - `solvent` and `vacuum` for solvation free energies, - or `complex` and `solvent` for binding free energies, + ``solvent`` and ``vacuum`` for solvation free energies, + or ``complex`` and ``solvent`` for binding free energies, each containing a list of NDArrays containing the corresponding - full system atom indices for each atom written in the production - trajectory files for each replica. + full system atom indices for each atom written in the PDB or + production trajectory files for each replica. """ indices: dict[str, list[Optional[npt.NDArray]]] = {} - for key in [self.bound_state, self.unbound_state]: + for key in [self.env_state, self.ref_state]: indices[key] = [] for pus in self.data[key].values(): # type: ignore[attr-defined] indices[key].append(pus[0].outputs["selection_indices"]) @@ -303,8 +309,8 @@ class AbsoluteSolvationProtocolResult(gufe.ProtocolResult, AbsoluteProtocolResul Protocol results with the output of a AbsoluteSolvationProtocol """ - bound_state = "solvent" - unbound_state = "vacuum" + env_state = "solvent" + ref_state = "vacuum" def get_individual_estimates(self) -> dict[str, list[tuple[Quantity, Quantity]]]: """ @@ -320,7 +326,7 @@ def get_individual_estimates(self) -> dict[str, list[tuple[Quantity, Quantity]]] """ dGs = {} - for state in [self.bound_state, self.unbound_state]: + for state in [self.env_state, self.ref_state]: state_dGs = [ (pus[0].outputs["unit_estimate"], pus[0].outputs["unit_estimate_error"]) for pus in self.data[state].values() @@ -389,8 +395,8 @@ class AbsoluteBindingProtocolResult(gufe.ProtocolResult, AbsoluteProtocolResultM Protocol results with the output of a AbsoluteBindingProtocol. """ - bound_state = "complex" - unbound_state = "solvent" + env_state = "complex" + ref_state = "solvent" def get_individual_estimates( self, diff --git a/src/openfe/protocols/openmm_rfe/__init__.py b/src/openfe/protocols/openmm_rfe/__init__.py index f0fe367c8..4083ba150 100644 --- a/src/openfe/protocols/openmm_rfe/__init__.py +++ b/src/openfe/protocols/openmm_rfe/__init__.py @@ -2,11 +2,35 @@ # For details, see https://github.com/OpenFreeEnergy/openfe from . import _rfe_utils -from .equil_rfe_settings import RelativeHybridTopologyProtocolSettings -from .hybridtop_protocol_results import RelativeHybridTopologyProtocolResult -from .hybridtop_protocols import RelativeHybridTopologyProtocol +from .equil_rfe_settings import ( + RBFEHTopProtocolSettings, + RelativeHybridTopologyProtocolSettings, + RHFEHTopProtocolSettings, +) +from .hybridtop_protocol_results import ( + RBFEHTopProtocolResult, + RelativeHybridTopologyProtocolResult, + RHFEHTopProtocolResult, +) +from .hybridtop_protocols import ( + RBFEHTopProtocol, + RelativeHybridTopologyProtocol, + RHFEHTopProtocol, +) from .hybridtop_units import ( HybridTopologyMultiStateAnalysisUnit, HybridTopologyMultiStateSimulationUnit, HybridTopologySetupUnit, + RBFEHTopComplexAnalysisUnit, + RBFEHTopComplexSetupUnit, + RBFEHTopComplexSimulationUnit, + RBFEHTopSolventAnalysisUnit, + RBFEHTopSolventSetupUnit, + RBFEHTopSolventSimulationUnit, + RHFEHTopSolventAnalysisUnit, + RHFEHTopSolventSetupUnit, + RHFEHTopSolventSimulationUnit, + RHFEHTopVacuumAnalysisUnit, + RHFEHTopVacuumSetupUnit, + RHFEHTopVacuumSimulationUnit, ) diff --git a/src/openfe/protocols/openmm_rfe/equil_rfe_settings.py b/src/openfe/protocols/openmm_rfe/equil_rfe_settings.py index 853df7173..5df213b21 100644 --- a/src/openfe/protocols/openmm_rfe/equil_rfe_settings.py +++ b/src/openfe/protocols/openmm_rfe/equil_rfe_settings.py @@ -102,11 +102,19 @@ class AlchemicalSettings(SettingsBaseModel): """ -class RelativeHybridTopologyProtocolSettings(Settings): +class BaseHTopProtocolSettings(SettingsBaseModel): + """ + Base configuration object for ``HTopProtocol`` and its subclasses. + + See Also + -------- + openfe.protocols.openmm_rfe.HTopProtocol + """ + protocol_repeats: int """ The number of completely independent repeats of the entire sampling - process. The mean of the repeats defines the final estimate of FE + process. The mean of the repeats defines the final estimate of the ΔΔG difference, while the variance between repeats is used as the uncertainty. """ @@ -117,28 +125,31 @@ def must_be_positive(cls, v): raise ValueError(errmsg) return v - # Inherited things - - forcefield_settings: OpenMMSystemGeneratorFFSettings - """Parameters to set up the force field with OpenMM Force Fields.""" thermo_settings: ThermoSettings """Settings for thermodynamic parameters.""" - # Things for creating the systems - solvation_settings: OpenMMSolvationSettings - """Settings for solvating the system.""" partial_charge_settings: OpenFFPartialChargeSettings """Settings for assigning partial charges to small molecules.""" # Alchemical settings - lambda_settings: LambdaSettings - """ - Lambda protocol settings including lambda windows and lambda functions. - """ alchemical_settings: AlchemicalSettings """ Alchemical protocol settings including soft core scaling. """ + + +class RelativeHybridTopologyProtocolSettings(BaseHTopProtocolSettings): + forcefield_settings: OpenMMSystemGeneratorFFSettings + """Parameters to set up the force field with OpenMM Force Fields.""" + + solvation_settings: OpenMMSolvationSettings + """Settings for solvating the system.""" + + lambda_settings: LambdaSettings + """ + Lambda protocol settings including lambda windows and lambda functions. + """ + simulation_settings: MultiStateSimulationSettings """ Settings for alchemical sampler. @@ -156,3 +167,126 @@ def must_be_positive(cls, v): """ Simulation output control settings. """ + + +class RBFEHTopProtocolSettings(BaseHTopProtocolSettings): + """ + Configuration object for ``RBFEHTopProtocol``. + + See Also + -------- + openfe.protocols.openmm_rfe.RBFEHTopProtocol + """ + + # Force field settings - only need one + forcefield_settings: OpenMMSystemGeneratorFFSettings + """Parameters to control assigning force field parameters using OpenMMForceFields.""" + + # Lambda schedule settings + solvent_lambda_settings: LambdaSettings + """ + Lambda protocol settings defining the lambda schedule, including + the number of lambda windows and scaling function for the solvent leg. + """ + complex_lambda_settings: LambdaSettings + """ + Lambda protocol settings defining the lambda schedule, including + the number of lambda windows and scaling function for the complex leg. + """ + + # Things for creating the systems + solvent_solvation_settings: OpenMMSolvationSettings + """Settings for solvating the solvent leg system.""" + complex_solvation_settings: OpenMMSolvationSettings + """Settings for solvating the complex leg system.""" + + # Simulation control settings + solvent_simulation_settings: MultiStateSimulationSettings + """ + Settings for controlling the solvent leg alchemical simulation. + """ + complex_simulation_settings: MultiStateSimulationSettings + """ + Settings for controlling the complex leg alchemical simulation. + """ + + # MD Engine things + engine_settings: OpenMMEngineSettings + """Settings specific to the OpenMM MD engine such as what compute platform to use.""" + + # Integrator control + solvent_integrator_settings: IntegratorSettings + """Settings for the solvent leg integrator such as timestep and barostat settings.""" + complex_integrator_settings: IntegratorSettings + """Settings for the complex leg integrator such as timestep and barostat settings.""" + + # Output control + solvent_output_settings: MultiStateOutputSettings + """ + Solvent leg simulation output (e.g. filenames) control settings. + """ + complex_output_settings: MultiStateOutputSettings + """ + Complex leg simulation output (e.g. filenames) control settings. + """ + + +class RHFEHTopProtocolSettings(BaseHTopProtocolSettings): + """ + Configuration object for ``RHFEHTopProtocol``. + + See Also + -------- + openfe.protocols.openmm_rfe.RHFEHTopProtocol + """ + + solvent_forcefield_settings: OpenMMSystemGeneratorFFSettings + """Parameters to control assigning force field parameters using OpenMMForceFields for the solvent leg.""" + vacuum_forcefield_settings: OpenMMSystemGeneratorFFSettings + """ + Parameters to control assigning the force field using OpenMMForceFields for the vacuum leg. + Must use a ``nonbonded_method`` of ``nocutoff``. + """ + + solvation_settings: OpenMMSolvationSettings + """Settings for solvating the solvent leg system. Ignored by the vacuum leg.""" + + solvent_lambda_settings: LambdaSettings + """ + Lambda protocol settings defining the lambda schedule, including + the number of lambda windows and scaling function for the solvent leg. + """ + vacuum_lambda_settings: LambdaSettings + """ + Lambda protocol settings defining the lambda schedule, including + the number of lambda windows and scaling function for the vacuum leg. + """ + + solvent_simulation_settings: MultiStateSimulationSettings + """ + Settings for controlling the solvent leg alchemical simulation. + """ + vacuum_simulation_settings: MultiStateSimulationSettings + """ + Settings for controlling the vacuum leg alchemical simulation. + """ + + # Engine settings control hardware usage + solvent_engine_settings: OpenMMEngineSettings + """Settings specific to the OpenMM MD engine for the solvent leg, such as what compute platform to use.""" + vacuum_engine_settings: OpenMMEngineSettings + """Settings specific to the OpenMM MD engine for the vacuum leg, such as what compute platform to use.""" + + solvent_integrator_settings: IntegratorSettings + """Settings for the solvent leg integrator such as timestep and barostat settings.""" + vacuum_integrator_settings: IntegratorSettings + """Settings for the vacuum leg integrator such as timestep settings.""" + + solvent_output_settings: MultiStateOutputSettings + """ + Solvent leg simulation output (e.g. filenames) control settings. + """ + vacuum_output_settings: MultiStateOutputSettings + """ + Vacuum leg simulation output (e.g. filenames) control settings. + """ diff --git a/src/openfe/protocols/openmm_rfe/hybridtop_protocol_results.py b/src/openfe/protocols/openmm_rfe/hybridtop_protocol_results.py index e67637d63..7fa6fa399 100644 --- a/src/openfe/protocols/openmm_rfe/hybridtop_protocol_results.py +++ b/src/openfe/protocols/openmm_rfe/hybridtop_protocol_results.py @@ -5,6 +5,7 @@ OpenMM and OpenMMTools in a Perses-like manner. """ +import itertools import logging import pathlib import warnings @@ -244,3 +245,367 @@ def production_iterations(self) -> list[float]: production_lengths = [pus[0].outputs["production_iterations"] for pus in self.data.values()] return production_lengths + + +class HTopProtocolResultMixin: + """ + Mixin providing the shared utilities for two-leg hybrid topology + ProtocolResults (``RBFEHTopProtocolResult``, ``RHFEHTopProtocolResult``). + + Subclasses must define the class attributes ``env_state`` and + ``ref_state``, naming the two legs of the thermodynamic cycle stored in + ``self.data`` (e.g. ``env_state = "complex"``, ``ref_state = "solvent"``). + """ + + env_state: str + ref_state: str + + def __init__(self, **data): + super().__init__(**data) + # data is a mapping of leg: str(repeat_id): list[protocolunitresults] + # TODO: Detect when we have extensions and stitch these together? + if any( + len(pur_list) > 2 + for pur_list in itertools.chain( + self.data[self.env_state].values(), self.data[self.ref_state].values() + ) + ): + raise NotImplementedError("Can't stitch together results yet") + + def get_individual_estimates(self) -> dict[str, list[tuple[Quantity, Quantity]]]: + """ + Get the individual estimate of the free energies for both legs. + + Returns + ------- + dGs : dict[str, list[tuple[openff.units.Quantity, openff.units.Quantity]]] + A dictionary, keyed for each leg of the thermodynamic cycle, e.g. + ``solvent`` and ``complex`` for a relaltive binding free energy or + ``solvent`` and ``vacuum`` for a relative hydration free energy, + with lists of tuples containing the individual free energy estimates + and associated MBAR uncertainties for each repeat of that simulation type. + """ + dGs = {} + + for state in [self.env_state, self.ref_state]: + dGs[state] = [ + (pus[0].outputs["unit_estimate"], pus[0].outputs["unit_estimate_error"]) + for pus in self.data[state].values() + ] + + return dGs + + @staticmethod + def _get_average(estimates: list[tuple[Quantity, Quantity]]) -> Quantity: + u = estimates[0][0].u + dGs = [i[0].to(u).m for i in estimates] + return np.average(dGs) * u + + @staticmethod + def _get_stdev(estimates: list[tuple[Quantity, Quantity]]) -> Quantity: + u = estimates[0][0].u + dGs = [i[0].to(u).m for i in estimates] + # use the unbiased sample standard deviation (ddof=1) as the repeats are sampled from the + # (inaccessible) population of possible repeats. + std = np.std(dGs, ddof=1) + if np.isnan(std): + std = 0.0 + return std * u + + def get_estimate(self) -> Quantity: + """Get the relative free energy estimate for this calculation. + + Returns + ------- + ddG : openff.units.Quantity + The difference free energy. This is a Quantity defined + with units. + """ + individual_estimates = self.get_individual_estimates() + env_dG = self._get_average(individual_estimates[self.env_state]) + ref_dG = self._get_average(individual_estimates[self.ref_state]) + + return env_dG - ref_dG + + def get_uncertainty(self) -> Quantity: + """Get the relative free energy error for this calculation. + + Returns + ------- + err : openff.units.Quantity + The unbiased standard deviation between estimates of the relative + free energy. This is a Quantity defined with units. + """ + individual_estimates = self.get_individual_estimates() + env_err = self._get_stdev(individual_estimates[self.env_state]) + ref_err = self._get_stdev(individual_estimates[self.ref_state]) + + return np.sqrt(env_err**2 + ref_err**2) + + def get_forward_and_reverse_energy_analysis( + self, + ) -> dict[str, list[Optional[dict[str, Union[npt.NDArray, Quantity]]]]]: + """ + Get the reverse and forward analysis of the free energies for both + legs of the thermodynamic cycle. + + Returns + ------- + forward_reverse : dict[str, list[Optional[dict[str, Union[npt.NDArray, openff.units.Quantity]]]]] + A dictionary, keyed by leg of the thermodynamic cycle, e.g. ``solvent`` + and ``vacuum`` for a relative hydration free energy or ``solvent`` and + ``complex`` for a relative binding free energy, with each + entry containing a list of dictionaries with the forward and + reverse analysis of each repeat of that simulation type. + + The forward and reverse analysis dictionaries contain: + - `fractions`: npt.NDArray + The fractions of data used for the estimates + - `forward_DGs`, `reverse_DGs`: openff.units.Quantity + The forward and reverse estimates for each fraction of data. + A fraction at which MBAR failed to converge is recorded as + ``NaN`` in both directions. + - `forward_dDGs`, `reverse_dDGs`: openff.units.Quantity + The forward and reverse estimate uncertainty for each + fraction of data (``NaN`` wherever the estimate is ``NaN``). + + A cycle leg list entry is ``None`` only when MBAR could not obtain + an estimate from the *full* set of uncorrelated samples (the + fraction 1.0 estimate, i.e. the reported free energy). If MBAR + fails only at a lower fraction, that fraction is recorded as + ``NaN`` (see ``forward_DGs`` above) and the remaining fractions + are retained, so the entry is still a dictionary. + + Raises + ------ + UserWarning + * If any of the forward and reverse dictionaries are ``None`` in a + given thermodynamic cycle leg. + """ + forward_reverse: dict[str, list[Optional[dict[str, Union[npt.NDArray, Quantity]]]]] = {} + + for key in [self.env_state, self.ref_state]: + forward_reverse[key] = [ + pus[0].outputs["forward_and_reverse_energies"] + for pus in self.data[key].values() # type: ignore[attr-defined] + ] + + if None in forward_reverse[key]: + wmsg = ( + "One or more ``None`` entries were found in the forward " + f"and reverse dictionaries of the repeats of the {key} " + "calculations. This indicates that MBAR could not obtain a " + "free energy estimate from the full set of uncorrelated " + "samples for that repeat." + ) + warnings.warn(wmsg) + + return forward_reverse + + def get_overlap_matrices(self) -> dict[str, list[dict[str, npt.NDArray]]]: + """ + Get the MBAR overlap estimates for both legs of the simulation. + + Returns + ------- + overlap_stats : dict[str, list[dict[str, npt.NDArray]]] + A dictionary keyed by leg of the thermodynamic cycle, e.g. + ``solvent`` and ``vacuum`` for a relative hydration free energy + or ``solvent`` and ``complex`` for a relative binding free energy, + with each entry containing a list of dictionaries with the MBAR overlap + estimates of each repeat of that simulation type. + + The underlying MBAR dictionaries contain the following keys: + * ``scalar``: One minus the largest nontrivial eigenvalue + * ``eigenvalues``: The sorted (descending) eigenvalues of the + overlap matrix + * ``matrix``: Estimated overlap matrix of observing a sample from + state i in state j + """ + # Loop through and get the repeats and get the matrices + overlap_stats: dict[str, list[dict[str, npt.NDArray]]] = {} + + for key in [self.env_state, self.ref_state]: + overlap_stats[key] = [ + pus[0].outputs["unit_mbar_overlap"] + for pus in self.data[key].values() # type: ignore[attr-defined] + ] + + return overlap_stats + + def get_replica_transition_statistics(self) -> dict[str, list[dict[str, npt.NDArray]]]: + """ + Get the replica exchange transition statistics for both legs of the + thermodynamic cycle. + + Note + ---- + This is currently only available in cases where a replica exchange + simulation was run. + + Returns + ------- + repex_stats : dict[str, list[dict[str, npt.NDArray]]] + A dictionary keyed by leg of the thermodynamic cycle, e.g. + ``solvent`` and ``vacuum`` for a relative hydration free energy or ``solvent`` and + ``complex`` for a relative binding free energy, with each + entry containing a list of dictionaries with the replica + transition statistics for each repeat of that simulation type. + + The replica transition statistics dictionaries contain the following: + * ``eigenvalues``: The sorted (descending) eigenvalues of the + lambda state transition matrix + * ``matrix``: The transition matrix estimate of a replica switching + from state i to state j. + """ + repex_stats: dict[str, list[dict[str, npt.NDArray]]] = {} + try: + for key in [self.env_state, self.ref_state]: + repex_stats[key] = [ + pus[0].outputs["replica_exchange_statistics"] + for pus in self.data[key].values() # type: ignore[attr-defined] + ] + except KeyError: + errmsg = "Replica exchange statistics were not found, did you run a repex calculation?" + raise ValueError(errmsg) + + return repex_stats + + def get_replica_states(self) -> dict[str, list[npt.NDArray]]: + """ + Get the timeseries of replica states for both simulation legs. + + Returns + ------- + replica_states : dict[str, list[npt.NDArray]] + Dictionary keyed by leg of the thermodynamic cycle, e.g. + ``solvent`` and ``vacuum`` for a relative hydration free energy or ``solvent`` and + ``complex`` for a relative binding free energy, with lists of + replica states timeseries for each repeat of that simulation type. + """ + replica_states: dict[str, list[npt.NDArray]] = { + self.env_state: [], + self.ref_state: [], + } + + def is_file(filename: str): + p = pathlib.Path(filename) + if not p.exists(): + errmsg = f"File could not be found {p}" + raise ValueError(errmsg) + return p + + def get_replica_state(nc, chk): + nc = is_file(nc) + dir_path = nc.parents[0] + chk = is_file(dir_path / chk).name + + reporter = multistate.MultiStateReporter( + storage=nc, checkpoint_storage=chk, open_mode="r" + ) + + retval = np.asarray(reporter.read_replica_thermodynamic_states()) + reporter.close() + + return retval + + for key in [self.env_state, self.ref_state]: + for pus in self.data[key].values(): # type: ignore[attr-defined] + states = get_replica_state( + pus[0].outputs["trajectory"], + pus[0].outputs["checkpoint"], + ) + replica_states[key].append(states) + + return replica_states + + def equilibration_iterations(self) -> dict[str, list[float]]: + """ + Returns the number of equilibration iterations for each repeat of + both legs of the calculation. + + Returns + ------- + equilibration_lengths : dict[str, list[float]] + Dictionary keyed for each leg of the thermodynamic cycle, e.g. + ``solvent`` and ``vacuum`` for a relative hydration free energy or + ``solvent`` and ``complex`` for a relative binding free energy, + with lists of the number of equilibration iterations for each + repeat of that simulation type. + """ + equilibration_lengths: dict[str, list[float]] = {} + + for key in [self.env_state, self.ref_state]: + equilibration_lengths[key] = [ + pus[0].outputs["equilibration_iterations"] + for pus in self.data[key].values() # type: ignore[attr-defined] + ] + + return equilibration_lengths + + def production_iterations(self) -> dict[str, list[float]]: + """ + Returns the number of uncorrelated production samples for each + repeat of both legs of the calculation. + + Returns + ------- + production_lengths : dict[str, list[float]] + Dictionary keyed for each leg of the thermodynamic cycle, e.g. + ``solvent`` and ``vacuum`` for a relative hydration free energy or + ``solvent`` and ``complex`` for a relative binding free energy, + with lists of the number of uncorrelated production samples for + each repeat of that simulation type. + """ + production_lengths: dict[str, list[float]] = {} + + for key in [self.env_state, self.ref_state]: + production_lengths[key] = [ + pus[0].outputs["production_iterations"] + for pus in self.data[key].values() # type: ignore[attr-defined] + ] + + return production_lengths + + def selection_indices(self) -> dict[str, list[Optional[npt.NDArray]]]: + """ + Get the system selection indices used to write PDB and trajectory + files, for both legs of the calculation. + + Returns + ------- + indices : dict[str, list[Optional[npt.NDArray]]] + A dictionary keyed by leg of the thermodynamic cycle, e.g. + ``solvent`` and ``vacuum`` for a relative hydration free energy or + ``solvent`` and ``complex`` for a relative binding free energy, + each containing a list of NDArrays with the corresponding full system + atom indices for each atom written in the PDB or production trajectory + files for each replica. + """ + indices: dict[str, list[Optional[npt.NDArray]]] = {} + + for key in [self.env_state, self.ref_state]: + indices[key] = [ + pus[0].outputs["selection_indices"] + for pus in self.data[key].values() # type: ignore[attr-defined] + ] + + return indices + + +class RBFEHTopProtocolResult(gufe.ProtocolResult, HTopProtocolResultMixin): + """ + Protocol results with the output of a ``RBFEHTopProtocol``. + """ + + env_state = "complex" + ref_state = "solvent" + + +class RHFEHTopProtocolResult(gufe.ProtocolResult, HTopProtocolResultMixin): + """ + Protocol results with the output of a ``RHFEHTopProtocol``. + """ + + env_state = "solvent" + ref_state = "vacuum" \ No newline at end of file diff --git a/src/openfe/protocols/openmm_rfe/hybridtop_protocols.py b/src/openfe/protocols/openmm_rfe/hybridtop_protocols.py index c615933b5..f2e3556cd 100644 --- a/src/openfe/protocols/openmm_rfe/hybridtop_protocols.py +++ b/src/openfe/protocols/openmm_rfe/hybridtop_protocols.py @@ -50,13 +50,31 @@ OpenFFPartialChargeSettings, OpenMMEngineSettings, OpenMMSolvationSettings, + RBFEHTopProtocolSettings, RelativeHybridTopologyProtocolSettings, + RHFEHTopProtocolSettings, +) +from .hybridtop_protocol_results import ( + RBFEHTopProtocolResult, + RelativeHybridTopologyProtocolResult, + RHFEHTopProtocolResult, ) -from .hybridtop_protocol_results import RelativeHybridTopologyProtocolResult from .hybridtop_units import ( HybridTopologyMultiStateAnalysisUnit, HybridTopologyMultiStateSimulationUnit, HybridTopologySetupUnit, + RBFEHTopComplexAnalysisUnit, + RBFEHTopComplexSetupUnit, + RBFEHTopComplexSimulationUnit, + RBFEHTopSolventAnalysisUnit, + RBFEHTopSolventSetupUnit, + RBFEHTopSolventSimulationUnit, + RHFEHTopSolventAnalysisUnit, + RHFEHTopSolventSetupUnit, + RHFEHTopSolventSimulationUnit, + RHFEHTopVacuumAnalysisUnit, + RHFEHTopVacuumSetupUnit, + RHFEHTopVacuumSimulationUnit, ) logger = logging.getLogger(__name__) @@ -84,138 +102,21 @@ ) -class RelativeHybridTopologyProtocol(gufe.Protocol): +class BaseHybridTopologyProtocol(gufe.Protocol): """ - Relative Free Energy calculations using a Hybrid Topology scheme - using OpenMM and OpenMMTools. - - Based on `Perses `_ - - See Also - -------- - :mod:`openfe.protocols` - :class:`openfe.protocols.openmm_rfe.RelativeHybridTopologySettings` - :class:`openfe.protocols.openmm_rfe.RelativeHybridTopologyResult` - :class:`openfe.protocols.openmm_rfe.RelativeHybridTopologyProtocolUnit` + Shared validation and DAG-construction logic for the hybrid topology + Protocols (``RelativeHybridTopologyProtocol``, ``RBFEHTopProtocol``, + ``RHFEHTopProtocol``). """ - result_cls = RelativeHybridTopologyProtocolResult - _settings_cls = RelativeHybridTopologyProtocolSettings - _settings: RelativeHybridTopologyProtocolSettings - - @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 RelativeHybridTopologyProtocolSettings( - protocol_repeats=3, - forcefield_settings=settings.OpenMMSystemGeneratorFFSettings(), - thermo_settings=settings.ThermoSettings( - temperature=298.15 * offunit.kelvin, - pressure=1 * offunit.bar, - ), - partial_charge_settings=OpenFFPartialChargeSettings(), - solvation_settings=OpenMMSolvationSettings(), - alchemical_settings=AlchemicalSettings(softcore_LJ="gapsys"), - lambda_settings=LambdaSettings(), - simulation_settings=MultiStateSimulationSettings( - equilibration_length=1.0 * offunit.nanosecond, - production_length=5.0 * offunit.nanosecond, - ), - engine_settings=OpenMMEngineSettings(), - integrator_settings=IntegratorSettings(), - output_settings=MultiStateOutputSettings(), - ) - - @classmethod - def _adaptive_settings( - cls, - stateA: ChemicalSystem, - stateB: ChemicalSystem, - mapping: gufe.LigandAtomMapping | list[gufe.LigandAtomMapping], - initial_settings: None | RelativeHybridTopologyProtocolSettings = None, - ) -> RelativeHybridTopologyProtocolSettings: - """ - Get the recommended OpenFE settings for this protocol based on the input states involved in the - transformation. - - These are intended as a suitable starting point for creating an instance of this protocol, which can be further - customized before performing a Protocol. - - Parameters - ---------- - stateA : ChemicalSystem - The initial state of the transformation. - stateB : ChemicalSystem - The final state of the transformation. - mapping : LigandAtomMapping | list[LigandAtomMapping] - The mapping(s) between transforming components in stateA and stateB. - initial_settings : None | RelativeHybridTopologyProtocolSettings, optional - Initial settings to base the adaptive settings on. If None, default settings are used. - - Returns - ------- - RelativeHybridTopologyProtocolSettings - The recommended settings for this protocol based on the input states. - - Notes - ----- - - If the transformation involves a change in net charge, the settings are adapted to use a more expensive - protocol with 22 lambda windows and 20 ns production length per window. - - If both states contain a ProteinComponent, the solvation padding is set to 1 nm. - - If initial_settings is provided, the adaptive settings are based on a copy of these settings. - """ - # use initial settings or default settings - # this is needed for the CLI so we don't override user settings - if initial_settings is not None: - protocol_settings = initial_settings.model_copy(deep=True) - else: - protocol_settings = cls.default_settings() - - if isinstance(mapping, list): - mapping = mapping[0] - - if mapping.get_alchemical_charge_difference() != 0: - # apply the recommended charge change settings taken from the industry benchmarking as fast settings not validated - # - info = ( - "Charge changing transformation between ligands " - f"{mapping.componentA.name} and {mapping.componentB.name}. " - "A more expensive protocol with 22 lambda windows, sampled " - "for 20 ns each, will be used here." - ) - logger.info(info) - protocol_settings.alchemical_settings.explicit_charge_correction = True - protocol_settings.simulation_settings.production_length = 20 * offunit.nanosecond - protocol_settings.simulation_settings.n_replicas = 22 - protocol_settings.lambda_settings.lambda_windows = 22 - - # adapt the solvation padding based on the system components - if stateA.contains(ProteinComponent): - protocol_settings.solvation_settings.solvent_padding = 1 * offunit.nanometer - - # adapt the barostat based on the system components - if stateA.contains(ProteinMembraneComponent): - protocol_settings.integrator_settings.barostat = "MonteCarloMembraneBarostat" - - return protocol_settings - @staticmethod - def _validate_endstates( + def _validate_endstate_alchemical_components( stateA: ChemicalSystem, stateB: ChemicalSystem, ) -> None: """ - Validates the end states for the RFE protocol. + Validates that there is exactly one alchemical + SmallMoleculeComponent per end state. Parameters ---------- @@ -494,129 +395,734 @@ def _validate_simulation_settings( simulation_settings=simulation_settings, ) - def _validate( + def _create_phased_units( self, stateA: ChemicalSystem, stateB: ChemicalSystem, - mapping: gufe.ComponentMapping | list[gufe.ComponentMapping] | None, - extends: gufe.ProtocolDAGResult | None = None, - ) -> None: - # Check we're not trying to extend - if extends: - # This technically should be NotImplementedError - # but gufe.Protocol.validate calls `_validate` wrapped around an - # except for NotImplementedError, so we can't raise it here - raise ValueError("Can't extend simulations yet") + mapping: Optional[Union[gufe.ComponentMapping, list[gufe.ComponentMapping]]], + phases: list[str], + unit_classes: dict[str, dict[str, type[gufe.ProtocolUnit]]], + label: str, + ) -> list[gufe.ProtocolUnit]: + """ + Shared helper to build setup/simulation/analysis units for each leg + of a multi-leg hybrid topology Protocol. - # Validate the end states - system_validation.validate_chemical_system(stateA) - system_validation.validate_chemical_system(stateB) - self._validate_endstates(stateA, stateB) + Parameters + ---------- + stateA, stateB : ChemicalSystem + The end states of the transformation. + mapping : ComponentMapping | list[ComponentMapping] | None + The mapping between transforming components. + phases : list[str] + The names of the legs to create units for, e.g. + ``["solvent", "complex"]``. + unit_classes : dict[str, dict[str, type[gufe.ProtocolUnit]]] + A dictionary, keyed by leg, of dictionaries mapping ``"setup"``, + ``"simulation"``, and ``"analysis"`` to the ProtocolUnit + subclasses to instantiate for that leg. + label : str + A short label used when naming the created units, e.g. + ``"RBFE HybridTopology"``. - # Validate the mapping + Returns + ------- + list[gufe.ProtocolUnit] + The flattened list of created units across all legs and repeats. + """ alchem_comps = system_validation.get_alchemical_components(stateA, stateB) - self._validate_mapping(mapping, alchem_comps) + ligandmapping = mapping[0] if isinstance(mapping, list) else mapping - # Validate the small molecule components - self._validate_smcs(stateA, stateB) + Anames = ",".join(c.name for c in alchem_comps["stateA"]) + Bnames = ",".join(c.name for c in alchem_comps["stateB"]) - # Validate solvent component - nonbond = self.settings.forcefield_settings.nonbonded_method - system_validation.validate_solvent(stateA, nonbond) + protocol_units: dict[str, list[gufe.ProtocolUnit]] = {phase: [] for phase in phases} - # Validate the BaseSolventComponents - base_solvent = stateA.get_components_of_type(BaseSolventComponent) - if len(base_solvent) > 1: - errmsg = "Multiple BaseSolventComponents found, only one is supported." - raise ValueError(errmsg) + for i in range(self.settings.protocol_repeats): + repeat_id = int(uuid.uuid4()) + for phase in phases: + setup = unit_classes[phase]["setup"]( + protocol=self, + stateA=stateA, + stateB=stateB, + ligandmapping=ligandmapping, + alchemical_components=alchem_comps, + generation=0, + repeat_id=repeat_id, + name=( + f"{label} Setup: {Anames} to {Bnames} {phase} leg: repeat {i} generation 0" + ), + ) - # Validate solvation settings - settings_validation.validate_openmm_solvation_settings(self.settings.solvation_settings) + simulation = unit_classes[phase]["simulation"]( + protocol=self, + setup_results=setup, + generation=0, + repeat_id=repeat_id, + name=( + f"{label} Simulation: {Anames} to {Bnames} {phase} leg: repeat {i} generation 0" + ), + ) - # Validate protein component - system_validation.validate_protein(stateA) + analysis = unit_classes[phase]["analysis"]( + protocol=self, + setup_results=setup, + simulation_results=simulation, + generation=0, + repeat_id=repeat_id, + name=( + f"{label} Analysis: {Anames} to {Bnames} {phase} leg: repeat {i} generation 0" + ), + ) - # Validate the barostat used in combination with the protein component - system_validation.validate_barostat(stateA, self.settings.integrator_settings.barostat) + protocol_units[phase] += [setup, simulation, analysis] - # Validate charge difference - # Note: validation depends on the mapping & solvent component checks - if stateA.contains(SolventComponent): - solv_comp = stateA.get_components_of_type(SolventComponent)[0] - elif stateA.contains(SolvatedPDBComponent): - solv_comp = stateA.get_components_of_type(SolvatedPDBComponent)[0] - else: - solv_comp = None + return [unit for phase in phases for unit in protocol_units[phase]] - self._validate_charge_difference( - mapping=mapping[0] if isinstance(mapping, list) else mapping, - nonbonded_method=self.settings.forcefield_settings.nonbonded_method, - explicit_charge_correction=self.settings.alchemical_settings.explicit_charge_correction, - solvent_component=solv_comp, - ) + def _gather_phased( + self, + protocol_dag_results: Iterable[gufe.ProtocolDAGResult], + phases: list[str], + ) -> dict[str, dict[str, list[gufe.ProtocolUnitResult]]]: + """ + Shared helper to gather Protocol results for a multi-leg hybrid + topology Protocol, grouping ``Analysis`` unit results first by leg + (via the ``simtype`` output) then by ``repeat_id``, sorted by + generation within each repeat. - # Validate integrator things - settings_validation.validate_timestep( - self.settings.forcefield_settings.hydrogen_mass, - self.settings.integrator_settings.timestep, - ) + Parameters + ---------- + protocol_dag_results : Iterable[gufe.ProtocolDAGResult] + The set of all ProtocolDAGResults to gather. + phases : list[str] + The names of the legs to gather results for, e.g. + ``["solvent", "complex"]``. - # Validate simulation & output settings - self._validate_simulation_settings( - self.settings.simulation_settings, - self.settings.integrator_settings, - self.settings.output_settings, - ) + Returns + ------- + dict[str, dict[str, list[gufe.ProtocolUnitResult]]] + A dictionary, keyed by leg, of dictionaries mapping + ``repeat_id`` to a sorted list of ``ProtocolUnitResult``. + """ + unsorted_repeats: dict[str, dict[Any, list[gufe.ProtocolUnitResult]]] = { + phase: defaultdict(list) for phase in phases + } - # Validate alchemical settings - # PR #125 temporarily pin lambda schedule spacing to n_replicas - if ( - self.settings.simulation_settings.n_replicas - != self.settings.lambda_settings.lambda_windows - ): - errmsg = ( - "Number of replicas in ``simulation_settings``: " - f"{self.settings.simulation_settings.n_replicas} must equal " - "the number of lambda windows in lambda_settings: " - f"{self.settings.lambda_settings.lambda_windows}." - ) - raise ValueError(errmsg) + for d in protocol_dag_results: + pu: gufe.ProtocolUnitResult + for pu in d.protocol_unit_results: + # We only need the analysis units that are ok + if ("Analysis" not in pu.name) or (not pu.ok()): + continue - def _create( - self, - stateA: ChemicalSystem, - stateB: ChemicalSystem, - mapping: Optional[Union[gufe.ComponentMapping, list[gufe.ComponentMapping]]], - extends: Optional[gufe.ProtocolDAGResult] = None, - ) -> list[gufe.ProtocolUnit]: - # validate inputs - self.validate(stateA=stateA, stateB=stateB, mapping=mapping, extends=extends) + phase = pu.outputs["simtype"] + unsorted_repeats[phase][pu.outputs["repeat_id"]].append(pu) - # get alchemical components and mapping - alchem_comps = system_validation.get_alchemical_components(stateA, stateB) - ligandmapping = mapping[0] if isinstance(mapping, list) else mapping + repeats: dict[str, dict[str, list[gufe.ProtocolUnitResult]]] = { + phase: {} for phase in phases + } + for phase in phases: + for k, v in unsorted_repeats[phase].items(): + repeats[phase][str(k)] = sorted(v, key=lambda x: x.outputs["generation"]) - # 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"]) + return repeats - # DAG dependency is setup -> simulation -> analysis - # |---------------------> - setup_units = [] - simulation_units = [] - analysis_units = [] - for i in range(self.settings.protocol_repeats): - repeat_id = int(uuid.uuid4()) +class RBFEHTopProtocol(BaseHybridTopologyProtocol): + """ + Relative Binding Free Energy calculations using a hybrid topology scheme + using OpenMM and OpenMMTools. - setup = HybridTopologySetupUnit( - protocol=self, - stateA=stateA, - stateB=stateB, - ligandmapping=ligandmapping, - alchemical_components=alchem_comps, - generation=0, - repeat_id=repeat_id, + Base on `Perses `_ + + See Also + -------- + :mod:`openfe.protocols` + # TODO - add more see alsos + """ + + result_cls = RBFEHTopProtocolResult + _settings_cls = RBFEHTopProtocolSettings + _settings: RBFEHTopProtocolSettings + + @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 RBFEHTopProtocolSettings( + protocol_repeats=3, + forcefield_settings=settings.OpenMMSystemGeneratorFFSettings(), + thermo_settings=settings.ThermoSettings( + temperature=298.15 * offunit.kelvin, + pressure=1 * offunit.bar, + ), + partial_charge_settings=OpenFFPartialChargeSettings(), + solvent_solvation_settings=OpenMMSolvationSettings(), + complex_solvation_settings=OpenMMSolvationSettings( + solvent_padding=1.0 * offunit.nanometer, + ), + alchemical_settings=AlchemicalSettings(softcore_LJ="gapsys"), + complex_lambda_settings=LambdaSettings(), + solvent_lambda_settings=LambdaSettings(), + solvent_simulation_settings=MultiStateSimulationSettings( + n_replicas=11, + equilibration_length=1.0 * offunit.nanosecond, + production_length=5.0 * offunit.nanosecond, + ), + complex_simulation_settings=MultiStateSimulationSettings( + n_replicas=11, + equilibration_length=1.0 * offunit.nanosecond, + production_length=5.0 * offunit.nanosecond, + ), + engine_settings=OpenMMEngineSettings(), + solvent_integrator_settings=IntegratorSettings(), + complex_integrator_settings=IntegratorSettings(), + solvent_output_settings=MultiStateOutputSettings( + output_structure="alchemical_system.pdb", + output_filename="solvent.nc", + checkpoint_storage_filename="solvent_checkpoint.nc", + ), + complex_output_settings=MultiStateOutputSettings( + output_structure="alchemical_system.pdb", + output_filename="complex.nc", + checkpoint_storage_filename="complex_checkpoint.nc", + ), + ) # fmt: skip + + @classmethod + def _adaptive_settings( + cls, + stateA: ChemicalSystem, + stateB: ChemicalSystem, + mapping: gufe.LigandAtomMapping | list[gufe.LigandAtomMapping], + initial_settings: None | RBFEHTopProtocolSettings = None, + ) -> RBFEHTopProtocolSettings: + """ + Get the recommended OpenFE settings for this protocol based on the input states involved in the + transformation. + + These are intended as a suitable starting point for creating an instance of this protocol, which can be further + customized before performing a Protocol. + + Parameters + ---------- + stateA : ChemicalSystem + The initial state of the transformation. + stateB : ChemicalSystem + The final state of the transformation. + mapping : LigandAtomMapping | list[LigandAtomMapping] + The mapping(s) between transforming components in stateA and stateB. + initial_settings : None | RBFEHTopProtocolSettings, optional + Initial settings to base the adaptive settings on. If None, default settings are used. + + Returns + ------- + RBFEHTopProtocolSettings + The recommended settings for this protocol based on the input states. + + Notes + ----- + - If the transformation involves a change in net charge, both legs' settings are adapted to + use a more expensive protocol with 22 lambda windows and 20 ns production length per window. + - If both states contain a ProteinComponent, the complex leg's solvation padding is set to 1 nm. + - If initial_settings is provided, the adaptive settings are based on a copy of these settings. + """ + if initial_settings is not None: + protocol_settings = initial_settings.model_copy(deep=True) + else: + protocol_settings = cls.default_settings() + + if isinstance(mapping, list): + mapping = mapping[0] + + if mapping.get_alchemical_charge_difference() != 0: + # apply the recommended charge change settings taken from the industry benchmarking as fast settings not validated + # + info = ( + "Charge changing transformation between ligands " + f"{mapping.componentA.name} and {mapping.componentB.name}. " + "A more expensive protocol with 22 lambda windows, sampled " + "for 20 ns each, will be used here for both legs." + ) + logger.info(info) + protocol_settings.alchemical_settings.explicit_charge_correction = True + for sim_settings in ( + protocol_settings.solvent_simulation_settings, + protocol_settings.complex_simulation_settings, + ): + sim_settings.production_length = 20 * offunit.nanosecond + sim_settings.n_replicas = 22 + for lambda_settings in ( + protocol_settings.solvent_lambda_settings, + protocol_settings.complex_lambda_settings, + ): + lambda_settings.lambda_windows = 22 + + # adapt the solvation padding based on the system components + if stateA.contains(ProteinComponent): + protocol_settings.complex_solvation_settings.solvent_padding = 1 * offunit.nanometer + + # adapt the barostat based on the system components + if stateA.contains(ProteinMembraneComponent): + protocol_settings.complex_integrator_settings.barostat = "MonteCarloMembraneBarostat" + + return protocol_settings + + @staticmethod + def _validate_endstates( + stateA: ChemicalSystem, + stateB: ChemicalSystem, + ) -> None: + """ + A complex relative transformation is defined (in terms of gufe components) + as starting from one or more ligands and a protein in solvent and + ending up in a state with one ligand that is different. + + Parameters + ---------- + stateA : ChemicalSystem + The chemical system of end state A. + stateB : ChemicalSystem + The chemical system of end state B. + + Raises + ------ + ValueError + If there is no SolventComponent and no ProteinComponent + in either stateA or stateB. + If there are no or more than one alchemical components in state A. + If there are no or more than one alchemical components in state B. + If there are any alchemical components that are not SmallMoleculeComponents. + """ + if not stateA.contains(ProteinComponent): + errmsg = "No ProteinComponent found in stateA" + raise ValueError(errmsg) + + if not stateB.contains(ProteinComponent): + errmsg = "No ProteinComponent found in stateB" + raise ValueError(errmsg) + + system_validation.validate_protein(stateA) + system_validation.validate_protein(stateB) + + if not stateA.contains(BaseSolventComponent): + errmsg = "No SolventComponent found in stateA" + raise ValueError(errmsg) + + if not stateB.contains(BaseSolventComponent): + errmsg = "No SolventComponent found in stateB" + raise ValueError(errmsg) + + BaseHybridTopologyProtocol._validate_endstate_alchemical_components(stateA, stateB) + + def _validate( + self, + stateA: ChemicalSystem, + stateB: ChemicalSystem, + mapping: gufe.ComponentMapping | list[gufe.ComponentMapping] | None, + extends: gufe.ProtocolDAGResult | None = None, + ) -> None: + # Check we're not trying to extend + if extends: + raise ValueError("Can't extend simulations yet") + + # Validate the end states + system_validation.validate_chemical_system(stateA) + system_validation.validate_chemical_system(stateB) + self._validate_endstates(stateA, stateB) + + # Validate the mapping + alchem_comps = system_validation.get_alchemical_components(stateA, stateB) + self._validate_mapping(mapping, alchem_comps) + + # Validate the small molecule components + self._validate_smcs(stateA, stateB) + + # Validate solvent component + nonbond = self.settings.forcefield_settings.nonbonded_method + system_validation.validate_solvent(stateA, nonbond) + + # Validate the BaseSolventComponents + base_solvent = stateA.get_components_of_type(BaseSolventComponent) + if len(base_solvent) > 1: + errmsg = "Multiple BaseSolventComponents found, only one is supported." + raise ValueError(errmsg) + + # Validate solvation settings + settings_validation.validate_openmm_solvation_settings( + self.settings.solvent_solvation_settings + ) + settings_validation.validate_openmm_solvation_settings( + self.settings.complex_solvation_settings + ) + + # Validate the barostat used in combination with the protein component + system_validation.validate_barostat( + stateA, self.settings.complex_integrator_settings.barostat + ) + + # Validate charge difference + # Note: validation depends on the mapping & solvent component checks + if stateA.contains(SolventComponent): + solv_comp = stateA.get_components_of_type(SolventComponent)[0] + elif stateA.contains(SolvatedPDBComponent): + solv_comp = stateA.get_components_of_type(SolvatedPDBComponent)[0] + else: + solv_comp = None + + self._validate_charge_difference( + mapping=mapping[0] if isinstance(mapping, list) else mapping, + nonbonded_method=nonbond, + explicit_charge_correction=self.settings.alchemical_settings.explicit_charge_correction, + solvent_component=solv_comp, + ) + + # Validate integrator things + settings_validation.validate_timestep( + self.settings.forcefield_settings.hydrogen_mass, + self.settings.solvent_integrator_settings.timestep, + ) + settings_validation.validate_timestep( + self.settings.forcefield_settings.hydrogen_mass, + self.settings.complex_integrator_settings.timestep, + ) + + # Validate simulation, output & lambda settings for both legs + for leg in ("solvent", "complex"): + sim_settings = getattr(self.settings, f"{leg}_simulation_settings") + integrator_settings = getattr(self.settings, f"{leg}_integrator_settings") + output_settings = getattr(self.settings, f"{leg}_output_settings") + lambda_settings = getattr(self.settings, f"{leg}_lambda_settings") + + self._validate_simulation_settings(sim_settings, integrator_settings, output_settings) + + # PR #125 temporarily pin lambda schedule spacing to n_replicas + if sim_settings.n_replicas != lambda_settings.lambda_windows: + errmsg = ( + "Number of replicas in " + f"``{leg}_simulation_settings``: {sim_settings.n_replicas} must equal " + f"the number of lambda windows in ``{leg}_lambda_settings``: " + f"{lambda_settings.lambda_windows}." + ) + raise ValueError(errmsg) + + def _create( + self, + stateA: ChemicalSystem, + stateB: ChemicalSystem, + mapping: Optional[Union[gufe.ComponentMapping, list[gufe.ComponentMapping]]], + extends: Optional[gufe.ProtocolDAGResult] = None, + ) -> list[gufe.ProtocolUnit]: + self.validate(stateA=stateA, stateB=stateB, mapping=mapping, extends=extends) + + unit_classes: dict[str, dict[str, type[gufe.ProtocolUnit]]] = { + "solvent": { + "setup": RBFEHTopSolventSetupUnit, + "simulation": RBFEHTopSolventSimulationUnit, + "analysis": RBFEHTopSolventAnalysisUnit, + }, + "complex": { + "setup": RBFEHTopComplexSetupUnit, + "simulation": RBFEHTopComplexSimulationUnit, + "analysis": RBFEHTopComplexAnalysisUnit, + }, + } + + return self._create_phased_units( + stateA=stateA, + stateB=stateB, + mapping=mapping, + phases=["solvent", "complex"], + unit_classes=unit_classes, + label="RBFE HybridTopology", + ) + + def _gather( + self, protocol_dag_results: Iterable[gufe.ProtocolDAGResult] + ) -> dict[str, dict[str, Any]]: + return self._gather_phased(protocol_dag_results, phases=["solvent", "complex"]) + + +class RelativeHybridTopologyProtocol(BaseHybridTopologyProtocol): + """ + Relative Free Energy calculations using a Hybrid Topology scheme + using OpenMM and OpenMMTools. + + Based on `Perses `_ + + See Also + -------- + :mod:`openfe.protocols` + :class:`openfe.protocols.openmm_rfe.RelativeHybridTopologySettings` + :class:`openfe.protocols.openmm_rfe.RelativeHybridTopologyResult` + :class:`openfe.protocols.openmm_rfe.RelativeHybridTopologyProtocolUnit` + """ + + result_cls = RelativeHybridTopologyProtocolResult + _settings_cls = RelativeHybridTopologyProtocolSettings + _settings: RelativeHybridTopologyProtocolSettings + + @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 RelativeHybridTopologyProtocolSettings( + protocol_repeats=3, + forcefield_settings=settings.OpenMMSystemGeneratorFFSettings(), + thermo_settings=settings.ThermoSettings( + temperature=298.15 * offunit.kelvin, + pressure=1 * offunit.bar, + ), + partial_charge_settings=OpenFFPartialChargeSettings(), + solvation_settings=OpenMMSolvationSettings(), + alchemical_settings=AlchemicalSettings(softcore_LJ="gapsys"), + lambda_settings=LambdaSettings(), + simulation_settings=MultiStateSimulationSettings( + equilibration_length=1.0 * offunit.nanosecond, + production_length=5.0 * offunit.nanosecond, + ), + engine_settings=OpenMMEngineSettings(), + integrator_settings=IntegratorSettings(), + output_settings=MultiStateOutputSettings(), + ) + + @classmethod + def _adaptive_settings( + cls, + stateA: ChemicalSystem, + stateB: ChemicalSystem, + mapping: gufe.LigandAtomMapping | list[gufe.LigandAtomMapping], + initial_settings: None | RelativeHybridTopologyProtocolSettings = None, + ) -> RelativeHybridTopologyProtocolSettings: + """ + Get the recommended OpenFE settings for this protocol based on the input states involved in the + transformation. + + These are intended as a suitable starting point for creating an instance of this protocol, which can be further + customized before performing a Protocol. + + Parameters + ---------- + stateA : ChemicalSystem + The initial state of the transformation. + stateB : ChemicalSystem + The final state of the transformation. + mapping : LigandAtomMapping | list[LigandAtomMapping] + The mapping(s) between transforming components in stateA and stateB. + initial_settings : None | RelativeHybridTopologyProtocolSettings, optional + Initial settings to base the adaptive settings on. If None, default settings are used. + + Returns + ------- + RelativeHybridTopologyProtocolSettings + The recommended settings for this protocol based on the input states. + + Notes + ----- + - If the transformation involves a change in net charge, the settings are adapted to use a more expensive + protocol with 22 lambda windows and 20 ns production length per window. + - If both states contain a ProteinComponent, the solvation padding is set to 1 nm. + - If initial_settings is provided, the adaptive settings are based on a copy of these settings. + """ + # use initial settings or default settings + # this is needed for the CLI so we don't override user settings + if initial_settings is not None: + protocol_settings = initial_settings.model_copy(deep=True) + else: + protocol_settings = cls.default_settings() + + if isinstance(mapping, list): + mapping = mapping[0] + + if mapping.get_alchemical_charge_difference() != 0: + # apply the recommended charge change settings taken from the industry benchmarking as fast settings not validated + # + info = ( + "Charge changing transformation between ligands " + f"{mapping.componentA.name} and {mapping.componentB.name}. " + "A more expensive protocol with 22 lambda windows, sampled " + "for 20 ns each, will be used here." + ) + logger.info(info) + protocol_settings.alchemical_settings.explicit_charge_correction = True + protocol_settings.simulation_settings.production_length = 20 * offunit.nanosecond + protocol_settings.simulation_settings.n_replicas = 22 + protocol_settings.lambda_settings.lambda_windows = 22 + + # adapt the solvation padding based on the system components + if stateA.contains(ProteinComponent): + protocol_settings.solvation_settings.solvent_padding = 1 * offunit.nanometer + + # adapt the barostat based on the system components + if stateA.contains(ProteinMembraneComponent): + protocol_settings.integrator_settings.barostat = "MonteCarloMembraneBarostat" + + return protocol_settings + + @staticmethod + def _validate_endstates( + stateA: ChemicalSystem, + stateB: ChemicalSystem, + ) -> None: + """ + Validates the end states for the RFE protocol. + + Parameters + ---------- + stateA : ChemicalSystem + The chemical system of end state A. + stateB : ChemicalSystem + The chemical system of end state B. + + Raises + ------ + ValueError + * If either state contains more than one unique Component. + * If unique components are not SmallMoleculeComponents. + """ + BaseHybridTopologyProtocol._validate_endstate_alchemical_components(stateA, stateB) + + def _validate( + self, + stateA: ChemicalSystem, + stateB: ChemicalSystem, + mapping: gufe.ComponentMapping | list[gufe.ComponentMapping] | None, + extends: gufe.ProtocolDAGResult | None = None, + ) -> None: + # Check we're not trying to extend + if extends: + # This technically should be NotImplementedError + # but gufe.Protocol.validate calls `_validate` wrapped around an + # except for NotImplementedError, so we can't raise it here + raise ValueError("Can't extend simulations yet") + + # Validate the end states + system_validation.validate_chemical_system(stateA) + system_validation.validate_chemical_system(stateB) + self._validate_endstates(stateA, stateB) + + # Validate the mapping + alchem_comps = system_validation.get_alchemical_components(stateA, stateB) + self._validate_mapping(mapping, alchem_comps) + + # Validate the small molecule components + self._validate_smcs(stateA, stateB) + + # Validate solvent component + nonbond = self.settings.forcefield_settings.nonbonded_method + system_validation.validate_solvent(stateA, nonbond) + + # Validate the BaseSolventComponents + base_solvent = stateA.get_components_of_type(BaseSolventComponent) + if len(base_solvent) > 1: + errmsg = "Multiple BaseSolventComponents found, only one is supported." + raise ValueError(errmsg) + + # Validate solvation settings + settings_validation.validate_openmm_solvation_settings(self.settings.solvation_settings) + + # Validate protein component + system_validation.validate_protein(stateA) + + # Validate the barostat used in combination with the protein component + system_validation.validate_barostat(stateA, self.settings.integrator_settings.barostat) + + # Validate charge difference + # Note: validation depends on the mapping & solvent component checks + if stateA.contains(SolventComponent): + solv_comp = stateA.get_components_of_type(SolventComponent)[0] + elif stateA.contains(SolvatedPDBComponent): + solv_comp = stateA.get_components_of_type(SolvatedPDBComponent)[0] + else: + solv_comp = None + + self._validate_charge_difference( + mapping=mapping[0] if isinstance(mapping, list) else mapping, + nonbonded_method=self.settings.forcefield_settings.nonbonded_method, + explicit_charge_correction=self.settings.alchemical_settings.explicit_charge_correction, + solvent_component=solv_comp, + ) + + # Validate integrator things + settings_validation.validate_timestep( + self.settings.forcefield_settings.hydrogen_mass, + self.settings.integrator_settings.timestep, + ) + + # Validate simulation & output settings + self._validate_simulation_settings( + self.settings.simulation_settings, + self.settings.integrator_settings, + self.settings.output_settings, + ) + + # Validate alchemical settings + # PR #125 temporarily pin lambda schedule spacing to n_replicas + if ( + self.settings.simulation_settings.n_replicas + != self.settings.lambda_settings.lambda_windows + ): + errmsg = ( + "Number of replicas in ``simulation_settings``: " + f"{self.settings.simulation_settings.n_replicas} must equal " + "the number of lambda windows in lambda_settings: " + f"{self.settings.lambda_settings.lambda_windows}." + ) + raise ValueError(errmsg) + + def _create( + self, + stateA: ChemicalSystem, + stateB: ChemicalSystem, + mapping: Optional[Union[gufe.ComponentMapping, list[gufe.ComponentMapping]]], + extends: Optional[gufe.ProtocolDAGResult] = None, + ) -> list[gufe.ProtocolUnit]: + # validate inputs + self.validate(stateA=stateA, stateB=stateB, mapping=mapping, extends=extends) + + # get alchemical components and mapping + alchem_comps = system_validation.get_alchemical_components(stateA, stateB) + ligandmapping = mapping[0] if isinstance(mapping, list) else mapping + + # 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"]) + + # DAG dependency is setup -> simulation -> analysis + # |---------------------> + setup_units = [] + simulation_units = [] + analysis_units = [] + + for i in range(self.settings.protocol_repeats): + repeat_id = int(uuid.uuid4()) + + setup = HybridTopologySetupUnit( + protocol=self, + stateA=stateA, + stateB=stateB, + ligandmapping=ligandmapping, + alchemical_components=alchem_comps, + generation=0, + repeat_id=repeat_id, name=(f"HybridTopology Setup: {Anames} to {Bnames} repeat {i} generation 0"), ) @@ -662,3 +1168,259 @@ def _gather(self, protocol_dag_results: Iterable[gufe.ProtocolDAGResult]) -> dic # returns a dict of repeat_id: sorted list of ProtocolUnitResult return repeats + + +class RHFEHTopProtocol(BaseHybridTopologyProtocol): + """ + Relative Hydration Free Energy calculations using a hybrid topology + scheme using OpenMM and OpenMMTools. + + Base on `Perses `_ + + See Also + -------- + :mod:`openfe.protocols` + """ + + result_cls = RHFEHTopProtocolResult + _settings_cls = RHFEHTopProtocolSettings + _settings: RHFEHTopProtocolSettings + + @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 RHFEHTopProtocolSettings( + protocol_repeats=3, + solvent_forcefield_settings=settings.OpenMMSystemGeneratorFFSettings(), + vacuum_forcefield_settings=settings.OpenMMSystemGeneratorFFSettings( + nonbonded_method="nocutoff", + ), + thermo_settings=settings.ThermoSettings( + temperature=298.15 * offunit.kelvin, + pressure=1 * offunit.bar, + ), + partial_charge_settings=OpenFFPartialChargeSettings(), + solvation_settings=OpenMMSolvationSettings(), + alchemical_settings=AlchemicalSettings(softcore_LJ="gapsys"), + solvent_lambda_settings=LambdaSettings(), + vacuum_lambda_settings=LambdaSettings(), + solvent_simulation_settings=MultiStateSimulationSettings( + n_replicas=11, + equilibration_length=1.0 * offunit.nanosecond, + production_length=5.0 * offunit.nanosecond, + ), + vacuum_simulation_settings=MultiStateSimulationSettings( + n_replicas=11, + equilibration_length=0.5 * offunit.nanosecond, + production_length=2.0 * offunit.nanosecond, + ), + solvent_engine_settings=OpenMMEngineSettings(), + vacuum_engine_settings=OpenMMEngineSettings(), + solvent_integrator_settings=IntegratorSettings(), + vacuum_integrator_settings=IntegratorSettings(), + solvent_output_settings=MultiStateOutputSettings( + output_structure="alchemical_system.pdb", + output_filename="solvent.nc", + checkpoint_storage_filename="solvent_checkpoint.nc", + ), + vacuum_output_settings=MultiStateOutputSettings( + output_structure="alchemical_system.pdb", + output_filename="vacuum.nc", + checkpoint_storage_filename="vacuum_checkpoint.nc", + ), + ) # fmt: skip + + @staticmethod + def _validate_no_charge_difference(mapping: LigandAtomMapping) -> None: + """ + Validates that there is no net charge difference between the end + states. + + RHFEHTopProtocol's vacuum leg has no solvent to draw an + alchemical water from, so (unlike RelativeHybridTopologyProtocol and + RBFEHTopProtocol) it cannot support an explicit charge + correction. Net charge changing transformations are therefore not + supported at all by this Protocol. + + Parameters + ---------- + mapping : LigandAtomMapping + The mapping between the transforming components. + + Raises + ------ + ValueError + If a change in net charge is detected. + """ + difference = mapping.get_alchemical_charge_difference() + + if abs(difference) != 0: + errmsg = ( + f"A charge difference of {difference} is observed " + "between the end states. RHFEHTopProtocol does not " + "support net charge changing transformations." + ) + raise ValueError(errmsg) + + @staticmethod + def _validate_endstates( + stateA: ChemicalSystem, + stateB: ChemicalSystem, + ) -> None: + """ + A relative hydration transformation is defined (in terms of gufe + components) as starting from one or more ligands in solvent and + ending up in a state with one ligand that is different. No protein + component is allowed. + + Parameters + ---------- + stateA : ChemicalSystem + The chemical system of end state A. + stateB : ChemicalSystem + The chemical system of end state B. + + Raises + ------ + ValueError + If either state contains a ProteinComponent. + If there is no SolventComponent in either stateA or stateB. + If there are no or more than one alchemical components in state A. + If there are no or more than one alchemical components in state B. + If there are any alchemical components that are not SmallMoleculeComponents. + """ + if stateA.contains(ProteinComponent) or stateB.contains(ProteinComponent): + errmsg = "Protein components are not allowed for RHFEHTopProtocol." + raise ValueError(errmsg) + + if not stateA.contains(BaseSolventComponent): + errmsg = "No SolventComponent found in stateA" + raise ValueError(errmsg) + + if not stateB.contains(BaseSolventComponent): + errmsg = "No SolventComponent found in stateB" + raise ValueError(errmsg) + + BaseHybridTopologyProtocol._validate_endstate_alchemical_components(stateA, stateB) + + def _validate( + self, + stateA: ChemicalSystem, + stateB: ChemicalSystem, + mapping: gufe.ComponentMapping | list[gufe.ComponentMapping] | None, + extends: gufe.ProtocolDAGResult | None = None, + ) -> None: + # Check we're not trying to extend + if extends: + raise ValueError("Can't extend simulations yet") + + # Validate the end states + system_validation.validate_chemical_system(stateA) + system_validation.validate_chemical_system(stateB) + self._validate_endstates(stateA, stateB) + + # Validate the mapping + alchem_comps = system_validation.get_alchemical_components(stateA, stateB) + self._validate_mapping(mapping, alchem_comps) + + # Validate the small molecule components + self._validate_smcs(stateA, stateB) + + # Validate that there is no net charge change (unsupported here) + self._validate_no_charge_difference(mapping[0] if isinstance(mapping, list) else mapping) + + # Validate solvent & vacuum nonbonded method compatibility + solvent_nonbonded_method = self.settings.solvent_forcefield_settings.nonbonded_method + vacuum_nonbonded_method = self.settings.vacuum_forcefield_settings.nonbonded_method + + system_validation.validate_solvent(stateA, solvent_nonbonded_method) + + if vacuum_nonbonded_method.lower() != "nocutoff": + errmsg = ( + "Only the nocutoff nonbonded_method is supported for the " + f"vacuum leg, {vacuum_nonbonded_method} was passed." + ) + raise ValueError(errmsg) + + # Validate the BaseSolventComponents + base_solvent = stateA.get_components_of_type(BaseSolventComponent) + if len(base_solvent) > 1: + errmsg = "Multiple BaseSolventComponents found, only one is supported." + raise ValueError(errmsg) + + # Validate solvation settings + settings_validation.validate_openmm_solvation_settings(self.settings.solvation_settings) + + # Validate integrator things + settings_validation.validate_timestep( + self.settings.solvent_forcefield_settings.hydrogen_mass, + self.settings.solvent_integrator_settings.timestep, + ) + settings_validation.validate_timestep( + self.settings.vacuum_forcefield_settings.hydrogen_mass, + self.settings.vacuum_integrator_settings.timestep, + ) + + # Validate simulation, output & lambda settings for both legs + for leg in ("solvent", "vacuum"): + sim_settings = getattr(self.settings, f"{leg}_simulation_settings") + integrator_settings = getattr(self.settings, f"{leg}_integrator_settings") + output_settings = getattr(self.settings, f"{leg}_output_settings") + lambda_settings = getattr(self.settings, f"{leg}_lambda_settings") + + self._validate_simulation_settings(sim_settings, integrator_settings, output_settings) + + if sim_settings.n_replicas != lambda_settings.lambda_windows: + errmsg = ( + "Number of replicas in " + f"``{leg}_simulation_settings``: {sim_settings.n_replicas} must equal " + f"the number of lambda windows in ``{leg}_lambda_settings``: " + f"{lambda_settings.lambda_windows}." + ) + raise ValueError(errmsg) + + def _create( + self, + stateA: ChemicalSystem, + stateB: ChemicalSystem, + mapping: Optional[Union[gufe.ComponentMapping, list[gufe.ComponentMapping]]], + extends: Optional[gufe.ProtocolDAGResult] = None, + ) -> list[gufe.ProtocolUnit]: + self.validate(stateA=stateA, stateB=stateB, mapping=mapping, extends=extends) + + unit_classes: dict[str, dict[str, type[gufe.ProtocolUnit]]] = { + "solvent": { + "setup": RHFEHTopSolventSetupUnit, + "simulation": RHFEHTopSolventSimulationUnit, + "analysis": RHFEHTopSolventAnalysisUnit, + }, + "vacuum": { + "setup": RHFEHTopVacuumSetupUnit, + "simulation": RHFEHTopVacuumSimulationUnit, + "analysis": RHFEHTopVacuumAnalysisUnit, + }, + } + + return self._create_phased_units( + stateA=stateA, + stateB=stateB, + mapping=mapping, + phases=["solvent", "vacuum"], + unit_classes=unit_classes, + label="RHFE HybridTopology", + ) + + def _gather( + self, protocol_dag_results: Iterable[gufe.ProtocolDAGResult] + ) -> dict[str, dict[str, Any]]: + return self._gather_phased(protocol_dag_results, phases=["solvent", "vacuum"]) diff --git a/src/openfe/protocols/openmm_rfe/hybridtop_units.py b/src/openfe/protocols/openmm_rfe/hybridtop_units.py index 4fbdbb9a1..13c8a028a 100644 --- a/src/openfe/protocols/openmm_rfe/hybridtop_units.py +++ b/src/openfe/protocols/openmm_rfe/hybridtop_units.py @@ -13,8 +13,6 @@ import logging import os import pathlib -import subprocess -from itertools import chain from typing import Any import gufe @@ -51,9 +49,6 @@ _set_offmol_metadata, _set_offmol_resname, ) -from openfe.protocols.openmm_utils.omm_settings import ( - BasePartialChargeSettings, -) from ...analysis import plotting from ...utils import log_system_probe, without_oechem_backend @@ -70,23 +65,24 @@ serialize, ) from . import _rfe_utils -from ._rfe_utils.relative import HybridTopologyFactory from .equil_rfe_settings import ( AlchemicalSettings, IntegratorSettings, - LambdaSettings, MultiStateOutputSettings, MultiStateSimulationSettings, OpenFFPartialChargeSettings, - OpenMMEngineSettings, OpenMMSolvationSettings, - RelativeHybridTopologyProtocolSettings, ) logger = logging.getLogger(__name__) class HybridTopologyUnitMixin: + #: Label identifying which leg of a multi-leg thermodynamic cycle this + #: unit belongs to (e.g. ``"solvent"``, ``"complex"``, ``"vacuum"``). + #: ``None`` for the single-leg ``RelativeHybridTopologyProtocol`` units. + simtype: str | None = None + def _prepare( self, verbose: bool, @@ -117,10 +113,7 @@ def _set_optional_path(basepath): self.scratch_basepath = _set_optional_path(scratch_basepath) self.shared_basepath = _set_optional_path(shared_basepath) - @staticmethod - def _get_settings( - settings: RelativeHybridTopologyProtocolSettings, - ) -> dict[str, SettingsBaseModel]: + def _get_settings(self) -> dict[str, SettingsBaseModel]: """ Get a dictionary of Protocol settings. @@ -131,22 +124,61 @@ def _get_settings( Notes ----- We return a dict so that we can duck type behaviour between phases. - For example subclasses may contain both `solvent` and `complex` - settings, using this approach we can extract the relevant entry - to the same key and pass it to other methods in a seamless manner. + Subclasses/mixins for multi-leg Protocols (e.g. ``solvent`` and + ``complex``) override this method to extract the relevant + leg-specific settings to the same dict keys, so that the rest of the + unit's code can be shared seamlessly across legs. + + This default implementation reads the flat, single-leg settings used + by ``RelativeHybridTopologyProtocol``. """ - protocol_settings: dict[str, SettingsBaseModel] = {} - protocol_settings["forcefield_settings"] = settings.forcefield_settings - protocol_settings["thermo_settings"] = settings.thermo_settings - protocol_settings["alchemical_settings"] = settings.alchemical_settings - protocol_settings["lambda_settings"] = settings.lambda_settings - protocol_settings["charge_settings"] = settings.partial_charge_settings - protocol_settings["solvation_settings"] = settings.solvation_settings - protocol_settings["simulation_settings"] = settings.simulation_settings - protocol_settings["output_settings"] = settings.output_settings - protocol_settings["integrator_settings"] = settings.integrator_settings - protocol_settings["engine_settings"] = settings.engine_settings - return protocol_settings + settings = self._inputs["protocol"].settings # type: ignore[attr-defined] + + return { + "forcefield_settings": settings.forcefield_settings, + "thermo_settings": settings.thermo_settings, + "alchemical_settings": settings.alchemical_settings, + "lambda_settings": settings.lambda_settings, + "charge_settings": settings.partial_charge_settings, + "solvation_settings": settings.solvation_settings, + "simulation_settings": settings.simulation_settings, + "output_settings": settings.output_settings, + "integrator_settings": settings.integrator_settings, + "engine_settings": settings.engine_settings, + } + + def _get_base_components( + self, + ) -> tuple[ + SolventComponent | None, ProteinComponent | None, dict[SmallMoleculeComponent, OFFMolecule] + ]: + """ + Get the solvent, protein and small molecule components directly from + the ``stateA``/``stateB`` inputs, with no leg-specific stripping. + + Returns + ------- + solv_comp : SolventComponent | None + The solvent component, if any. + protein_comp : ProteinComponent | None + The protein component, if any. + small_mols : dict[SmallMoleculeComponent, openff.toolkit.Molecule] + Dictionary of small molecule components paired with their + OpenFF Molecule. + """ + stateA = self._inputs["stateA"] # type: ignore[attr-defined] + stateB = self._inputs["stateB"] # type: ignore[attr-defined] + + solvent_comp, protein_comp, smcs_A = system_validation.get_components(stateA) + _, _, smcs_B = system_validation.get_components(stateB) + + small_mols = {m: m.to_openff() for m in set(smcs_A).union(set(smcs_B))} + + # If there is a SolvatedPDBComponent, we set the solvent_comp + if isinstance(protein_comp, SolvatedPDBComponent): + solvent_comp = protein_comp + + return solvent_comp, protein_comp, small_mols @staticmethod def _verify_execution_environment( @@ -169,45 +201,39 @@ def _verify_execution_environment( raise ProtocolUnitExecutionError(errmsg) -class HybridTopologySetupUnit(gufe.ProtocolUnit, HybridTopologyUnitMixin): +class BaseHybridTopologySetupUnit(gufe.ProtocolUnit, HybridTopologyUnitMixin): """ - Setup unit for Hybrid Topology Protocol transformations. + Base setup unit for Hybrid Topology Protocol transformations. + + Subclasses (in combination with a components mixin, e.g. + ``HybridTopologyComplexComponentsMixin``) must provide a + ``_get_components(self)`` method. """ - @staticmethod def _get_components( - stateA: ChemicalSystem, stateB: ChemicalSystem - ) -> tuple[SolventComponent, ProteinComponent, dict[SmallMoleculeComponent, OFFMolecule]]: + self, + ) -> tuple[ + SolventComponent | None, ProteinComponent | None, dict[SmallMoleculeComponent, OFFMolecule] + ]: """ Get the components from the ChemicalSystem inputs. - Parameters - ---------- - stateA : ChemicalSystem - ChemicalSystem defining the state A components. - stateB : CHemicalSystem - ChemicalSystem defining the state B components. - Returns ------- - solv_comp : SolventComponent - The solvent component. - protein_comp : ProteinComponent - The protein component. + solv_comp : SolventComponent | None + The solvent component, if any. + protein_comp : ProteinComponent | None + The protein component, if any. small_mols : dict[SmallMoleculeComponent, openff.toolkit.Molecule] Dictionary of small molecule components paired with their OpenFF Molecule. - """ - solvent_comp, protein_comp, smcs_A = system_validation.get_components(stateA) - _, _, smcs_B = system_validation.get_components(stateB) - - small_mols = {m: m.to_openff() for m in set(smcs_A).union(set(smcs_B))} - # If there is a SolvatedPDBComponent, we set the solvent_comp - if isinstance(protein_comp, SolvatedPDBComponent): - solvent_comp = protein_comp - - return solvent_comp, protein_comp, small_mols + Note + ---- + Must be implemented (directly, or via a components mixin) in the + child class. + """ + raise NotImplementedError() @staticmethod def _assign_partial_charges( @@ -749,14 +775,14 @@ def run( self.logger.info("Starting system setup unit") # Get settings - settings = self._get_settings(self._inputs["protocol"].settings) + settings = self._get_settings() # Get components stateA = self._inputs["stateA"] stateB = self._inputs["stateB"] mapping = self._inputs["ligandmapping"] alchem_comps = self._inputs["alchemical_components"] - solvent_comp, protein_comp, small_mols = self._get_components(stateA, stateB) + solvent_comp, protein_comp, small_mols = self._get_components() alchemical = set(alchem_comps["stateA"]) | set(alchem_comps["stateB"]) @@ -880,6 +906,7 @@ def _execute( return { "repeat_id": self._inputs["repeat_id"], "generation": self._inputs["generation"], + "simtype": self.simtype, "openmm_version": openmm.__version__, "openfe_version": openfe.__version__, "gufe_version": gufe.__version__, @@ -887,9 +914,9 @@ def _execute( } -class HybridTopologyMultiStateSimulationUnit(gufe.ProtocolUnit, HybridTopologyUnitMixin): +class BaseHybridTopologyMultiStateSimulationUnit(gufe.ProtocolUnit, HybridTopologyUnitMixin): """ - Multi-state simulation (e.g. multi replica methods like hamiltonian + Base multi-state simulation (e.g. multi replica methods like hamiltonian replica exchange) unit for Hybrid Topology Protocol transformations. """ @@ -1360,7 +1387,7 @@ def run( self.logger.info("Starting simulation unit") # Get the settings - settings = self._get_settings(self._inputs["protocol"].settings) + settings = self._get_settings() # Check for a restart self.restart = self._check_restart( @@ -1489,13 +1516,14 @@ def _execute( return { "repeat_id": self._inputs["repeat_id"], "generation": self._inputs["generation"], + "simtype": self.simtype, **outputs, } -class HybridTopologyMultiStateAnalysisUnit(gufe.ProtocolUnit, HybridTopologyUnitMixin): +class BaseHybridTopologyMultiStateAnalysisUnit(gufe.ProtocolUnit, HybridTopologyUnitMixin): """ - Analysis unit for multi-state Hybrid Topology Protocol transformations. + Base analysis unit for multi-state Hybrid Topology Protocol transformations. """ @staticmethod @@ -1672,7 +1700,7 @@ def run( self.logger.info("Starting simulation analysis unit") # Get the settings - settings = self._get_settings(self._inputs["protocol"].settings) + settings = self._get_settings() # Energies analysis if verbose: @@ -1733,6 +1761,7 @@ def _execute( return { "repeat_id": self._inputs["repeat_id"], "generation": self._inputs["generation"], + "simtype": self.simtype, # We include various other outputs here to make # things easier when gathering. "pdb_structure": pdb_file, @@ -1741,3 +1770,244 @@ def _execute( "selection_indices": selection_indices, **outputs, } + + +class HybridTopologyComplexComponentsMixin: + """ + Components mixin returning the full set of components (protein, solvent, + small molecules) present in the ``stateA``/``stateB`` inputs, with no + leg-specific stripping. Used both by the single-leg + ``RelativeHybridTopologyProtocol`` units and by the complex leg of + ``RBFEHTopProtocol``. + """ + + def _get_components( + self, + ) -> tuple[ + SolventComponent | None, ProteinComponent | None, dict[SmallMoleculeComponent, OFFMolecule] + ]: + return self._get_base_components() + + +class HybridTopologySolventComponentsMixin: + """ + Components mixin for the solvent leg of a multi-leg Protocol: reuses + the full components, but always nulls out any protein component (be it + absent from the input already, as for ``RHFEHTopProtocol``, or + present and stripped, as for the solvent leg derived from + ``RBFEHTopProtocol``'s complex-shaped ChemicalSystem inputs). + """ + + def _get_components( + self, + ) -> tuple[SolventComponent | None, None, dict[SmallMoleculeComponent, OFFMolecule]]: + solvent_comp, _protein_comp, small_mols = self._get_base_components() # type: ignore[attr-defined] + return solvent_comp, None, small_mols + + +class HybridTopologyVacuumComponentsMixin: + """ + Components mixin for the vacuum leg of ``RHFEHTopProtocol``: reuses + the full components, but always nulls out the solvent component. + """ + + def _get_components( + self, + ) -> tuple[None, ProteinComponent | None, dict[SmallMoleculeComponent, OFFMolecule]]: + _solvent_comp, protein_comp, small_mols = self._get_base_components() # type: ignore[attr-defined] + return None, protein_comp, small_mols + + +class HybridTopologySetupUnit(HybridTopologyComplexComponentsMixin, BaseHybridTopologySetupUnit): + """ + Setup unit for Hybrid Topology Protocol transformations. + """ + + +class HybridTopologyMultiStateSimulationUnit(BaseHybridTopologyMultiStateSimulationUnit): + """ + Multi-state simulation (e.g. multi replica methods like hamiltonian + replica exchange) unit for Hybrid Topology Protocol transformations. + """ + + +class HybridTopologyMultiStateAnalysisUnit(BaseHybridTopologyMultiStateAnalysisUnit): + """ + Analysis unit for multi-state Hybrid Topology Protocol transformations. + """ + + +class RBFEComplexSettingsMixin: + """Settings mixin for the complex leg of ``RBFEHTopProtocol``.""" + + def _get_settings(self) -> dict[str, SettingsBaseModel]: + settings = self._inputs["protocol"].settings # type: ignore[attr-defined] + + return { + "forcefield_settings": settings.forcefield_settings, + "thermo_settings": settings.thermo_settings, + "alchemical_settings": settings.alchemical_settings, + "lambda_settings": settings.complex_lambda_settings, + "charge_settings": settings.partial_charge_settings, + "solvation_settings": settings.complex_solvation_settings, + "simulation_settings": settings.complex_simulation_settings, + "output_settings": settings.complex_output_settings, + "integrator_settings": settings.complex_integrator_settings, + "engine_settings": settings.engine_settings, + } + + +class RBFESolventSettingsMixin: + """Settings mixin for the solvent leg of ``RBFEHTopProtocol``.""" + + def _get_settings(self) -> dict[str, SettingsBaseModel]: + settings = self._inputs["protocol"].settings + + return { + "forcefield_settings": settings.forcefield_settings, + "thermo_settings": settings.thermo_settings, + "alchemical_settings": settings.alchemical_settings, + "lambda_settings": settings.solvent_lambda_settings, + "charge_settings": settings.partial_charge_settings, + "solvation_settings": settings.solvent_solvation_settings, + "simulation_settings": settings.solvent_simulation_settings, + "output_settings": settings.solvent_output_settings, + "integrator_settings": settings.solvent_integrator_settings, + "engine_settings": settings.engine_settings, + } + + +class RBFEHTopComplexSetupUnit( + HybridTopologyComplexComponentsMixin, RBFEComplexSettingsMixin, BaseHybridTopologySetupUnit +): + """Setup unit for the complex leg of ``RBFEHTopProtocol``.""" + + simtype = "complex" + + +class RBFEHTopComplexSimulationUnit( + RBFEComplexSettingsMixin, BaseHybridTopologyMultiStateSimulationUnit +): + """Multi-state simulation unit for the complex leg of ``RBFEHTopProtocol``.""" + + simtype = "complex" + + +class RBFEHTopComplexAnalysisUnit( + RBFEComplexSettingsMixin, BaseHybridTopologyMultiStateAnalysisUnit +): + """Analysis unit for the complex leg of ``RBFEHTopProtocol``.""" + + simtype = "complex" + + +class RBFEHTopSolventSetupUnit( + HybridTopologySolventComponentsMixin, RBFESolventSettingsMixin, BaseHybridTopologySetupUnit +): + """Setup unit for the solvent leg of ``RBFEHTopProtocol``.""" + + simtype = "solvent" + + +class RBFEHTopSolventSimulationUnit( + RBFESolventSettingsMixin, BaseHybridTopologyMultiStateSimulationUnit +): + """Multi-state simulation unit for the solvent leg of ``RBFEHTopProtocol``.""" + + simtype = "solvent" + + +class RBFEHTopSolventAnalysisUnit( + RBFESolventSettingsMixin, BaseHybridTopologyMultiStateAnalysisUnit +): + """Analysis unit for the solvent leg of ``RBFEHTopProtocol``.""" + + simtype = "solvent" + + +class RHFESolventSettingsMixin: + """Settings mixin for the solvent leg of ``RHFEHTopProtocol``.""" + + def _get_settings(self) -> dict[str, SettingsBaseModel]: + settings = self._inputs["protocol"].settings + + return { + "forcefield_settings": settings.solvent_forcefield_settings, + "thermo_settings": settings.thermo_settings, + "alchemical_settings": settings.alchemical_settings, + "lambda_settings": settings.solvent_lambda_settings, + "charge_settings": settings.partial_charge_settings, + "solvation_settings": settings.solvation_settings, + "simulation_settings": settings.solvent_simulation_settings, + "output_settings": settings.solvent_output_settings, + "integrator_settings": settings.solvent_integrator_settings, + "engine_settings": settings.solvent_engine_settings, + } + + +class RHFEVacuumSettingsMixin: + """Settings mixin for the vacuum leg of ``RHFEHTopProtocol``.""" + + def _get_settings(self) -> dict[str, SettingsBaseModel]: + settings = self._inputs["protocol"].settings # type: ignore[attr-defined] + + return { + "forcefield_settings": settings.vacuum_forcefield_settings, + "thermo_settings": settings.thermo_settings, + "alchemical_settings": settings.alchemical_settings, + "lambda_settings": settings.vacuum_lambda_settings, + "charge_settings": settings.partial_charge_settings, + # Solvation settings are ignored by the vacuum leg (no solvent to + # add), but are included here for a consistent settings shape. + "solvation_settings": settings.solvation_settings, + "simulation_settings": settings.vacuum_simulation_settings, + "output_settings": settings.vacuum_output_settings, + "integrator_settings": settings.vacuum_integrator_settings, + "engine_settings": settings.vacuum_engine_settings, + } + + +class RHFEHTopSolventSetupUnit( + HybridTopologySolventComponentsMixin, RHFESolventSettingsMixin, BaseHybridTopologySetupUnit +): + """Setup unit for the solvent leg of ``RHFEHTopProtocol``.""" + + simtype = "solvent" + + +class RHFEHTopSolventSimulationUnit( + RHFESolventSettingsMixin, BaseHybridTopologyMultiStateSimulationUnit +): + """Multi-state simulation unit for the solvent leg of ``RHFEHTopProtocol``.""" + + simtype = "solvent" + + +class RHFEHTopSolventAnalysisUnit( + RHFESolventSettingsMixin, BaseHybridTopologyMultiStateAnalysisUnit +): + """Analysis unit for the solvent leg of ``RHFEHTopProtocol``.""" + + simtype = "solvent" + + +class RHFEHTopVacuumSetupUnit( + HybridTopologyVacuumComponentsMixin, RHFEVacuumSettingsMixin, BaseHybridTopologySetupUnit +): + """Setup unit for the vacuum leg of ``RHFEHTopProtocol``.""" + + simtype = "vacuum" + + +class RHFEHTopVacuumSimulationUnit( + RHFEVacuumSettingsMixin, BaseHybridTopologyMultiStateSimulationUnit +): + """Multi-state simulation unit for the vacuum leg of ``RHFEHTopProtocol``.""" + + simtype = "vacuum" + + +class RHFEHTopVacuumAnalysisUnit(RHFEVacuumSettingsMixin, BaseHybridTopologyMultiStateAnalysisUnit): + """Analysis unit for the vacuum leg of ``RHFEHTopProtocol``.""" + + simtype = "vacuum"