From 1a3e6943d95acd6c32c2b4a44f8850be6514d758 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Fri, 19 Jun 2026 15:07:33 +0200 Subject: [PATCH 1/4] SepTop Boresch solvent prototype --- .../openmm_septop/equil_septop_method.py | 9 +- .../protocols/openmm_septop/septop_units.py | 270 ++++++------ .../openmm_septop/solvent_boresch.py | 235 +++++++++++ .../restraint_utils/geometry/boresch/dummy.py | 259 ++++++++++++ .../restraint_utils/openmm/omm_dummy.py | 296 +++++++++++++ .../openmm_septop/test_septop_protocol.py | 19 +- .../test_septop_solvent_restraints.py | 393 ++++++++++++++++++ .../restraints/test_dummy_boresch.py | 296 +++++++++++++ 8 files changed, 1620 insertions(+), 157 deletions(-) create mode 100644 src/openfe/protocols/openmm_septop/solvent_boresch.py create mode 100644 src/openfe/protocols/restraint_utils/geometry/boresch/dummy.py create mode 100644 src/openfe/protocols/restraint_utils/openmm/omm_dummy.py create mode 100644 src/openfe/tests/protocols/openmm_septop/test_septop_solvent_restraints.py create mode 100644 src/openfe/tests/protocols/restraints/test_dummy_boresch.py diff --git a/src/openfe/protocols/openmm_septop/equil_septop_method.py b/src/openfe/protocols/openmm_septop/equil_septop_method.py index 05d84a1bf..245ce42ca 100644 --- a/src/openfe/protocols/openmm_septop/equil_septop_method.py +++ b/src/openfe/protocols/openmm_septop/equil_septop_method.py @@ -68,10 +68,7 @@ ) from ..openmm_utils import settings_validation, system_validation -from ..restraint_utils.settings import ( - BoreschRestraintSettings, - DistanceRestraintSettings, -) +from ..restraint_utils.settings import BoreschRestraintSettings from .septop_protocol_results import SepTopProtocolResult from .septop_units import ( SepTopComplexAnalysisUnit, @@ -263,9 +260,7 @@ def _default_settings(cls): output_filename="complex.nc", checkpoint_storage_filename="complex_checkpoint.nc", ), - solvent_restraint_settings=DistanceRestraintSettings( - spring_constant=1000.0 * offunit.kilojoule_per_mole / offunit.nanometer**2, - ), + solvent_restraint_settings=BoreschRestraintSettings(), complex_restraint_settings=BoreschRestraintSettings(), ) # fmt: skip diff --git a/src/openfe/protocols/openmm_septop/septop_units.py b/src/openfe/protocols/openmm_septop/septop_units.py index 6ad261837..1ed90520c 100644 --- a/src/openfe/protocols/openmm_septop/septop_units.py +++ b/src/openfe/protocols/openmm_septop/septop_units.py @@ -39,10 +39,7 @@ from openfe.protocols.restraint_utils import geometry from openfe.protocols.restraint_utils.geometry.boresch import BoreschRestraintGeometry from openfe.protocols.restraint_utils.openmm import omm_restraints -from openfe.protocols.restraint_utils.openmm.omm_restraints import ( - BoreschRestraint, - add_force_in_separate_group, -) +from openfe.protocols.restraint_utils.openmm.omm_restraints import BoreschRestraint from ..openmm_utils import ( settings_validation, @@ -51,7 +48,6 @@ from ..openmm_utils.mdtraj_utils import mdtraj_from_openmm from ..restraint_utils.settings import ( BoreschRestraintSettings, - DistanceRestraintSettings, ) from .base_units import ( BaseSepTopAnalysisUnit, @@ -59,10 +55,41 @@ BaseSepTopSetupUnit, _pre_equilibrate, ) +from .solvent_boresch import add_solvent_boresch_restraints logger = logging.getLogger(__name__) +def _add_dummy_atoms_to_topology( + topology: openmm.app.Topology, + n_dummies: int = 6, +) -> None: + """ + Extend *topology* in-place with *n_dummies* dummy atoms. + + Each dummy is added as a single-atom residue named ``DUM`` in a new + chain, with element ``None`` and atom name ``DUM``. This ensures + the topology particle count matches the OpenMM System after + ``add_dummy_atoms_to_system`` has been called. + + Modifies the topology in-place (OpenMM Topology is mutable). + A deepcopy is deliberately avoided because it breaks the internal + atom-object identity that ``PDBFile.writeFooter`` relies on when + building CONECT records. + + Parameters + ---------- + topology : openmm.app.Topology + The topology to extend. Modified in-place. + n_dummies : int + Number of dummy atoms to append. Default 6 (3 per ligand). + """ + chain = topology.addChain() + for i in range(n_dummies): + residue = topology.addResidue("DUM", chain) + topology.addAtom(f"DUM{i}", None, residue) + + class SepTopComplexMixin: """ A mixin to get the components and the settings for the Complex Units. @@ -880,136 +907,80 @@ class SepTopSolventSetupUnit(SepTopSolventMixin, BaseSepTopSetupUnit): simtype = "solvent" - @staticmethod - def _update_positions( - mol_A: SmallMoleculeComponent, - mol_B: SmallMoleculeComponent, - ) -> SmallMoleculeComponent: - """ - Computes the amount to offset the second ligand by in the solution - phase during RBFE calculations and applies the offset to the ligand, - returning the SmallMoleculeComponent with the updated positions. - - Parameters - ---------- - mol_A: SmallMoleculeComponent - The SmallMoleculeComponent of ligand A - mol_B: SmallMoleculeComponent - The SmallMoleculeComponent of ligand B - Returns - ------- - updated_mol_B: SmallMoleculeComponent - The SmallMoleculeComponent of ligand B after updating its positions - to be a certain distance away from ligand A - """ - - # Convert SmallMolecule to Rdkit Molecule - rdmol_A = mol_A.to_rdkit() - rdmol_B = mol_B.to_rdkit() - # Offset ligand B from ligand A in the solvent - pos_ligandA = rdmol_A.GetConformers()[0].GetPositions() - pos_ligandB = rdmol_B.GetConformers()[0].GetPositions() - - ligand_1_radius = np.linalg.norm(pos_ligandA - pos_ligandA.mean(axis=0), axis=1).max() - ligand_2_radius = np.linalg.norm(pos_ligandB - pos_ligandB.mean(axis=0), axis=1).max() - ligand_distance = (ligand_1_radius + ligand_2_radius) * 1.5 - - ligand_offset = pos_ligandA.mean(0) - pos_ligandB.mean(0) - ligand_offset[0] += ligand_distance - - # Offset the ligandB. - pos_ligandB += ligand_offset - - # Extract updated system positions. - rdmol_B.GetConformers()[0].SetPositions(pos_ligandB) - - updated_mol_B = SmallMoleculeComponent(rdmol_B) - - return updated_mol_B - def _add_restraints( - self, - system: openmm.System, - ligand_1: Chem.rdchem.Mol, - ligand_2: Chem.rdchem.Mol, - ligand_1_inxs: list[int], - ligand_2_inxs: list[int], - settings: dict[str, SettingsBaseModel], - positions_AB: openmm.unit.Quantity, + self, + system: openmm.System, + rdmol_A: Chem.rdchem.Mol, + rdmol_B: Chem.rdchem.Mol, + ligand_A_idxs: list[int], + ligand_B_idxs: list[int], + settings: dict[str, SettingsBaseModel], + positions_AB: np.ndarray, ) -> tuple[ + Quantity, Quantity, openmm.System, + np.ndarray, + BoreschRestraintGeometry, + BoreschRestraintGeometry, ]: """ - Apply the distance restraint between the ligands. + Apply Boresch restraints for both ligands using analytically-placed + dummy atoms as host anchors. Parameters ---------- system: openmm.System - The OpenMM system where the restraints will be applied to. - ligand_1: Chem.rdchem.Mol - The RDKit Molecule of ligand A - ligand_2: Chem.rdchem.Mol - The RDKit Molecule of ligand B - ligand_1_idxs: list[int] - Atom indices from the ligand A in the system. - ligand_2_idxs: list[int] - Atom indices from the ligand B in the system. + The alchemical OpenMM system. Modified in-place. + rdmol_A: Chem.rdchem.Mol + Sanitised RDKit molecule for ligand A. + rdmol_B: Chem.rdchem.Mol + Sanitised RDKit molecule for ligand B. + ligand_A_idxs: list[int] + Atom indices of ligand A in the full system. + ligand_B_idxs: list[int] + Atom indices of ligand B in the full system. settings: dict[str, SettingsBaseModel] - The settings dict - positions_AB: openmm.unit.Quantity - The positions of the OpenMM system + Protocol settings dict, must contain ``"restraint_settings"`` + (a ``BoreschRestraintSettings``) and ``"thermo_settings"``. + positions_AB: np.ndarray + Full-system positions in Angstroms, shape ``(N, 3)``. Returns ------- - correction: unit.Quantity - Standard state correction for the harmonic distance restraint. + correction_A: Quantity + Boresch standard-state correction for ligand A (kJ/mol). + correction_B: Quantity + Boresch standard-state correction for ligand B (kJ/mol). system: openmm.System - The OpenMM system with the added restraints forces + The system with restraint forces added and dummy atoms appended. + positions_AB: np.ndarray + Extended positions array of shape ``(N + 6, 3)`` in Angstroms, + with dummy atom coordinates appended. + restraint_geom_A: BoreschRestraintGeometry + The restraint Geometry object for ligand A. ``host_atoms`` are + the 3 dummy atom indices anchoring ligand A. + restraint_geom_B: BoreschRestraintGeometry + The restraint Geometry object for ligand B. ``host_atoms`` are + the 3 dummy atom indices anchoring ligand B. """ - - if isinstance(settings["restraint_settings"], DistanceRestraintSettings): - rest_geom = geometry.harmonic.get_molecule_centers_restraint( - molA_rdmol=ligand_1, - molB_rdmol=ligand_2, - molA_idxs=ligand_1_inxs, - molB_idxs=ligand_2_inxs, - ) - - else: - # TODO turn this into a direction for different restraint types supported? - raise NotImplementedError("Other restraint types are not yet available") - - if self.verbose: - self.logger.info(f"restraint geometry is: {rest_geom}") - - distance = np.linalg.norm( - positions_AB[rest_geom.guest_atoms[0]] - positions_AB[rest_geom.host_atoms[0]] + corr_A, corr_B, system, positions_AB, geom_A, geom_B = add_solvent_boresch_restraints( + system=system, + positions_ang=positions_AB, + rdmol_A=rdmol_A, + rdmol_B=rdmol_B, + ligand_A_idxs=ligand_A_idxs, + ligand_B_idxs=ligand_B_idxs, + settings=settings, ) - k_distance = to_openmm(settings["restraint_settings"].spring_constant) - - force = openmm.HarmonicBondForce() - force.addBond( - rest_geom.guest_atoms[0], - rest_geom.host_atoms[0], - distance * openmm.unit.nanometers, - k_distance, - ) - force.setName("alignment_restraint") - # Add force to a separate force group - add_force_in_separate_group(system, force) - - # No correction necessary as only a single harmonic bond is applied between the ligands - correction = ( - from_openmm( - openmm.unit.MOLAR_GAS_CONSTANT_R - * to_openmm(settings["thermo_settings"].temperature) + if self.verbose: + self.logger.info( + f"Applied dummy-atom Boresch restraints for solvent leg. " + f"Standard state corrections: A={corr_A:.3f}, B={corr_B:.3f}" ) - * 0.0 - ) - return correction, system + return corr_A, corr_B, system, positions_AB, geom_A, geom_B def run( self, dry=False, verbose=True, scratch_basepath=None, shared_basepath=None @@ -1052,21 +1023,21 @@ def run( # 3. Assign partial charges self._assign_partial_charges(settings["charge_settings"], smc_comps_AB) - - # 4. Update the positions of ligand B: - # - solvent: Offset ligand B with respect to ligand A - smc_B = self._update_positions( - alchem_comps["stateA"][0], - alchem_comps["stateB"][0], - ) - smc_off_B = {smc_B: smc_B.to_openff()} + # + # # 4. Update the positions of ligand B: + # # - solvent: Offset ligand B with respect to ligand A + # smc_B = self._update_positions( + # alchem_comps["stateA"][0], + # alchem_comps["stateB"][0], + # ) + # smc_off_B = {smc_B: smc_B.to_openff()} # 5. Get the OpenMM systems omm_system_AB, omm_topology_AB, positions_AB, modeller_AB, comp_resids_AB = ( self.get_system( solv_comp, prot_comp, - smc_comps_A | smc_off_B, + smc_comps_A | smc_comps_B, settings, ) ) # fmt: skip @@ -1075,7 +1046,7 @@ def run( # system AB comp_atomids_AB = self._get_atom_indices(omm_topology_AB, comp_resids_AB) atom_indices_AB_A = comp_atomids_AB[alchem_comps["stateA"][0]] - atom_indices_AB_B = comp_atomids_AB[smc_B] + atom_indices_AB_B = comp_atomids_AB[alchem_comps["stateB"][0]] # 7. Create the alchemical system self.logger.info("Creating the alchemical system and applying restraints") @@ -1089,24 +1060,38 @@ def run( # 8. Apply Restraints rdmol_A = alchem_comps["stateA"][0].to_rdkit() - rdmol_B = smc_B.to_rdkit() + rdmol_B = alchem_comps["stateB"][0].to_rdkit() Chem.SanitizeMol(rdmol_A) Chem.SanitizeMol(rdmol_B) - corr, system = self._add_restraints( - alchemical_system, - rdmol_A, - rdmol_B, - atom_indices_AB_A, - atom_indices_AB_B, - settings, - positions_AB, + # positions_AB is extended by 6 rows (3 dummies per ligand) + positions_AB_ang = np.array( + positions_AB.value_in_unit(openmm.unit.angstrom)) + corr_A, corr_B, system, positions_AB_ang, restraint_geom_A, restraint_geom_B = ( + self._add_restraints( + alchemical_system, + rdmol_A, + rdmol_B, + atom_indices_AB_A, + atom_indices_AB_B, + settings, + positions_AB_ang, + ) ) + positions_AB = positions_AB_ang * openmm.unit.angstrom + + # Extend omm_topology_AB with 6 dummy atoms so the PDB captures all + # N+6 particles in the system. The run unit reads positions from this + # PDB, so the particle count must match the serialised system exactly. + _add_dummy_atoms_to_topology(omm_topology_AB, n_dummies=6) - # Write the full system PDB + # Write the full system PDB. + # positions_AB now has 6 extra rows for the dummy atoms, which are + # not in omm_topology_AB. Slice back to the original atom count. topology_file = self.shared_basepath / "topology.pdb" openmm.app.pdbfile.PDBFile.writeFile( - omm_topology_AB, positions_AB, open(topology_file, "w") + omm_topology_AB, positions_AB, + open(topology_file, "w") ) # Subselect system based on user inputs & write initial subsampled PDB @@ -1130,13 +1115,13 @@ def run( return { "system": system_outfile, "topology": topology_file, - "standard_state_correction": corr.to("kilocalorie_per_mole"), + "standard_state_correction_A": corr_A.to("kilocalorie_per_mole"), + "standard_state_correction_B": corr_B.to("kilocalorie_per_mole"), "selection_indices": selection_indices, "subsampled_pdb_structure": sub_pdb_structure, } else: return { - # Add in various objects we can used to test the system "system": system_outfile, "topology": topology_file, "system_AB": omm_system_AB, @@ -1144,6 +1129,8 @@ def run( "alchem_system": alchemical_system, "alchem_factory": alchemical_factory, "positions": positions_AB, + "restraint_geometry_A": restraint_geom_A, + "restraint_geometry_B": restraint_geom_B, "selection_indices": selection_indices, "subsampled_pdb_structure": sub_pdb_structure, } @@ -1172,14 +1159,15 @@ def _get_lambda_schedule( lambda_vdw_A = [1 - x for x in lambda_vdw_A] lambda_elec_B = [1 - x for x in lambda_elec_B] lambda_vdw_B = [1 - x for x in lambda_vdw_B] - # # Set lambda restraint for the solvent to 1 - # lambda_restraints = len(lambda_elec_A) * [1] lambdas["lambda_electrostatics_A"] = lambda_elec_A lambdas["lambda_sterics_A"] = lambda_vdw_A lambdas["lambda_electrostatics_B"] = lambda_elec_B lambdas["lambda_sterics_B"] = lambda_vdw_B - # lambdas['lambda_restraints'] = lambda_restraints + # Dummy-atom Boresch restraints are always fully on in the solvent leg + n_windows = len(lambda_elec_A) + lambdas["lambda_restraints_A"] = [1.0] * n_windows + lambdas["lambda_restraints_B"] = [1.0] * n_windows return lambdas diff --git a/src/openfe/protocols/openmm_septop/solvent_boresch.py b/src/openfe/protocols/openmm_septop/solvent_boresch.py new file mode 100644 index 000000000..fba84f769 --- /dev/null +++ b/src/openfe/protocols/openmm_septop/solvent_boresch.py @@ -0,0 +1,235 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe +""" +Solvent-leg Boresch restraint wiring for SepTopSolventSetupUnit. + +This module provides a drop-in replacement for the ``_add_restraints`` method +of ``SepTopSolventSetupUnit``. It replaces the single harmonic distance +restraint with two independent Boresch restraints (one per ligand), each +anchored to three dedicated dummy atoms. + +Lambda schedule +--------------- +The restraints use controlling parameter names +``"lambda_restraints_A"`` and ``"lambda_restraints_B"``. For the solvent +leg, both should remain at 1.0 throughout all lambda windows (i.e. the +restraints are always on). + +Standard state correction +------------------------- +For the solvent leg the Boresch correction has the *same* sign for both +ligands — they are both being restrained (neither is being released). +The net contribution to the RBFE therefore cancels when taking the +difference. We return it for bookkeeping but it should not be applied +asymmetrically as in the complex leg. +""" +from __future__ import annotations + +import numpy as np +import openmm +import openmm.unit as omm_unit +import MDAnalysis as mda +from MDAnalysis.coordinates.memory import MemoryReader +from gufe.settings.models import SettingsBaseModel +from openff.units import Quantity +from openff.units.openmm import to_openmm +from openmmtools.states import ThermodynamicState +from rdkit import Chem + +from openfe.protocols.restraint_utils import geometry +from openfe.protocols.restraint_utils.geometry.boresch import BoreschRestraintGeometry +from openfe.protocols.restraint_utils.geometry.boresch import find_guest_atom_candidates +from openfe.protocols.restraint_utils.openmm.omm_restraints import ( + BoreschRestraint, + add_force_in_separate_group, +) +from openfe.protocols.restraint_utils.settings import BoreschRestraintSettings + +from openfe.protocols.restraint_utils.geometry.boresch.dummy import find_dummy_atom_positions +from openfe.protocols.restraint_utils.openmm.omm_dummy import add_dummy_atoms_to_system + + +def add_solvent_boresch_restraints( + system: openmm.System, + positions_ang: np.ndarray, + rdmol_A: Chem.Mol, + rdmol_B: Chem.Mol, + ligand_A_idxs: list[int], + ligand_B_idxs: list[int], + settings: dict[str, SettingsBaseModel], +) -> tuple[ + Quantity, + Quantity, + openmm.System, + np.ndarray, + BoreschRestraintGeometry, + BoreschRestraintGeometry, +]: + """ + Add Boresch restraints for both ligands to the solvent-leg System, + using analytically-placed dummy atoms as hosts. + + Parameters + ---------- + system: + The (alchemical) OpenMM System to modify. **Modified in-place.** + positions_ang: + Full-system positions in Angstroms, shape ``(N, 3)``. Must already + reflect the desired input conformer positions for both ligands. + rdmol_A, rdmol_B: + Sanitised RDKit molecules for ligands A and B. + ligand_A_idxs, ligand_B_idxs: + Atom indices for each ligand in the full system. + settings: + The protocol settings dict, must contain: + * ``"restraint_settings"`` — a ``BoreschRestraintSettings`` instance + * ``"thermo_settings"`` — for temperature / pressure + + Returns + ------- + correction_A: + Boresch standard-state correction for ligand A (kJ/mol). + correction_B: + Boresch standard-state correction for ligand B (kJ/mol). + system: + The modified System (same object, returned for convenience). + positions_ang: + Extended positions array with dummy atom coordinates appended, + shape ``(N + 6, 3)``. + geom_A: + The BoreschRestraintGeometry applied to ligand A (host_atoms are + the 3 dummy atom indices for ligand A; guest_atoms are G0/G1/G2). + geom_B: + The BoreschRestraintGeometry applied to ligand B. + + Raises + ------ + TypeError + If ``settings["restraint_settings"]`` is not a + ``BoreschRestraintSettings`` instance. + ValueError + If no suitable ligand anchor atoms can be found for either ligand. + """ + restraint_settings = settings["restraint_settings"] + if not isinstance(restraint_settings, BoreschRestraintSettings): + raise TypeError( + "Solvent-leg dummy Boresch restraints require a " + f"BoreschRestraintSettings instance, got " + f"{type(restraint_settings).__name__}." + ) + + # 1. Add 6 dummy atoms to the system (3 per ligand). + # Positions are initialised to zero; we fill them in below. + system, positions_ang, dummy_idxs_A = add_dummy_atoms_to_system( + system, positions_ang, n_dummies=3 + ) + system, positions_ang, dummy_idxs_B = add_dummy_atoms_to_system( + system, positions_ang, n_dummies=3 + ) + + # 2. For each ligand: find guest anchor atoms (G0/G1/G2), place the + # 3 dummy atoms analytically around them, then build the + # BoreschRestraintGeometry by delegating to the same + # find_boresch_restraint entry point used by the complex leg, + # via its guest_restraint_atoms_idxs / host_restraint_atoms_idxs + # override path (we already know exactly which atoms to use, + # so we skip its host/guest search logic). + def _build_geometry_for_ligand( + rdmol: Chem.Mol, + lig_idxs: list[int], + dummy_idxs: list[int], + ) -> BoreschRestraintGeometry: + """ + Find guest anchor atoms, place dummies, write dummy positions into + positions_ang in-place, and return a BoreschRestraintGeometry. + """ + n_atoms = positions_ang.shape[0] + u = mda.Universe.empty(n_atoms, trajectory=True) + u.load_new(positions_ang[np.newaxis, :, :], format=MemoryReader) + + anchors = find_guest_atom_candidates( + universe=u, + rdmol=rdmol, + guest_idxs=lig_idxs, + rmsf_cutoff=restraint_settings.rmsf_cutoff, + ) + if not anchors: + raise ValueError( + "No suitable ligand anchor atoms found for dummy Boresch " + "restraint. Try using a ligand with aromatic rings or more " + "rigid heavy atoms." + ) + g0_idx, g1_idx, g2_idx = anchors[0] + + # Place the 3 dummy atoms analytically and write their positions + # into positions_ang (and the live MDA universe) before measuring + # the restraint geometry. + p_d0, p_d1, p_d2 = find_dummy_atom_positions( + positions_ang[g0_idx], + positions_ang[g1_idx], + positions_ang[g2_idx], + ) + positions_ang[dummy_idxs[0]] = p_d0 + positions_ang[dummy_idxs[1]] = p_d1 + positions_ang[dummy_idxs[2]] = p_d2 + u.atoms.positions = positions_ang + + # guest_restraint_atoms_idxs are positions *within* lig_idxs, not + # absolute system indices (mirrors how find_boresch_restraint's + # override path indexes into universe.atoms[guest_idxs]). + guest_restraint_atoms_idxs = [ + lig_idxs.index(g0_idx), + lig_idxs.index(g1_idx), + lig_idxs.index(g2_idx), + ] + # host_restraint_atoms_idxs select all 3 dummies, in placement + # order (D0, D1, D2), since dummy_idxs is exactly the host pool. + host_restraint_atoms_idxs = [0, 1, 2] + + return geometry.boresch.find_boresch_restraint( + universe=u, + guest_rdmol=rdmol, + guest_idxs=lig_idxs, + host_idxs=dummy_idxs, + guest_restraint_atoms_idxs=guest_restraint_atoms_idxs, + host_restraint_atoms_idxs=host_restraint_atoms_idxs, + ) + + geom_A = _build_geometry_for_ligand(rdmol_A, ligand_A_idxs, dummy_idxs_A) + geom_B = _build_geometry_for_ligand(rdmol_B, ligand_B_idxs, dummy_idxs_B) + + # 4. Add the Boresch forces via the existing BoreschRestraint class. + restraint_A = BoreschRestraint(restraint_settings) + restraint_B = BoreschRestraint(restraint_settings) + + thermodynamic_state = ThermodynamicState( + system, + temperature=to_openmm(settings["thermo_settings"].temperature), + pressure=to_openmm(settings["thermo_settings"].pressure), + ) + + restraint_A.add_force( + thermodynamic_state, + geom_A, + controlling_parameter_name="lambda_restraints_A", + ) + restraint_B.add_force( + thermodynamic_state, + geom_B, + controlling_parameter_name="lambda_restraints_B", + ) + + # 5. Standard state corrections. + # In the solvent leg both ligands are restrained throughout, so + # corrections are equal in magnitude. We return them for + # bookkeeping; they cancel in the RBFE cycle. + correction_A = restraint_A.get_standard_state_correction( + thermodynamic_state, geom_A + ) + correction_B = restraint_B.get_standard_state_correction( + thermodynamic_state, geom_B + ) + + system = thermodynamic_state.get_system(remove_thermostat=True) + + return correction_A, correction_B, system, positions_ang, geom_A, geom_B \ No newline at end of file diff --git a/src/openfe/protocols/restraint_utils/geometry/boresch/dummy.py b/src/openfe/protocols/restraint_utils/geometry/boresch/dummy.py new file mode 100644 index 000000000..f16af84ee --- /dev/null +++ b/src/openfe/protocols/restraint_utils/geometry/boresch/dummy.py @@ -0,0 +1,259 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe +""" +Dummy-atom Boresch restraint geometry for the solvent leg of SepTop. + +In the solvent leg there is no host molecule to anchor Boresch restraints to. +Instead, three dummy atoms are placed analytically around each ligand's input +conformer so that all Boresch angles and dihedrals are well-defined (not near +0 or 180 degrees). The dummies carry no nonbonded interactions and have a +very large mass, making them effectively immobile throughout the simulation. + +Geometry construction +--------------------- +Given three ligand anchor atoms G0, G1, G2 at positions p0, p1, p2: + + D2 G2 + - - + - - + D1 - - D0 -- G0 - - G1 + +D0 is placed along the normal n to the G0-G1-G2 plane, at distance r0 from +G0. This guarantees theta_B (D0-G0-G1) = 90 degrees. + +D1 is placed so that theta_A (D1-D0-G0) = 90 degrees, in the plane spanned +by n and (p1 - p0). + +D2 is placed so that phi_A (D2-D1-D0-G0) = 60 degrees. + +All angles are validated to be away from the singular values 0 and 180 +degrees. +""" +from __future__ import annotations + +import warnings + +import numpy as np + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +#: Distance (Å) from G0 to D0. +_DUMMY_BOND_LENGTH_A: float = 5.0 + +#: Minimum safe angle (radians) away from 0 or pi. +_ANGLE_WARN_THRESHOLD_RAD: float = np.deg2rad(10.0) + + +# --------------------------------------------------------------------------- +# Geometry helpers +# --------------------------------------------------------------------------- + + +def _normalise(v: np.ndarray) -> np.ndarray: + """Return the unit vector of *v*. Raises if the norm is zero.""" + n = np.linalg.norm(v) + if n < 1e-8: + raise ValueError(f"Cannot normalise a near-zero vector: {v}") + return v / n + + +def _perpendicular_in_plane( + a: np.ndarray, + b: np.ndarray, +) -> np.ndarray: + """ + Return a unit vector that is perpendicular to *a* and lies in the plane + spanned by *a* and *b*. + + Parameters + ---------- + a: + The primary direction (will be normalised). + b: + A second vector that, together with *a*, defines the plane. + + Returns + ------- + np.ndarray + Unit vector perpendicular to *a* in the (a, b) plane. + """ + a_hat = _normalise(a) + # Remove the a component from b + b_perp = b - np.dot(b, a_hat) * a_hat + return _normalise(b_perp) + + +def _rotate_around_axis( + v: np.ndarray, + axis: np.ndarray, + angle_rad: float, +) -> np.ndarray: + """ + Rotate vector *v* around *axis* by *angle_rad* using Rodrigues' formula. + """ + axis = _normalise(axis) + return ( + v * np.cos(angle_rad) + + np.cross(axis, v) * np.sin(angle_rad) + + axis * np.dot(axis, v) * (1 - np.cos(angle_rad)) + ) + + +def _check_angle_safe(angle_rad: float, name: str) -> None: + """ + Warn if *angle_rad* is within ``_ANGLE_WARN_THRESHOLD_RAD`` of 0 or pi. + """ + if angle_rad < _ANGLE_WARN_THRESHOLD_RAD or angle_rad > np.pi - _ANGLE_WARN_THRESHOLD_RAD: + warnings.warn( + f"Boresch angle {name} = {np.degrees(angle_rad):.1f} deg is close " + "to a singular value (0 or 180 deg). Consider choosing different " + "ligand anchor atoms.", + UserWarning, + stacklevel=2, + ) + + +# --------------------------------------------------------------------------- +# Core dummy placement +# --------------------------------------------------------------------------- + + +def find_dummy_atom_positions( + p_g0: np.ndarray, + p_g1: np.ndarray, + p_g2: np.ndarray, + bond_length_a: float = _DUMMY_BOND_LENGTH_A, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Analytically place three dummy atoms (D0, D1, D2) around three ligand + anchor atoms (G0, G1, G2) such that the resulting Boresch angles and + dihedrals are well-defined. + + All input positions must be in **Angstroms**. + + Parameters + ---------- + p_g0, p_g1, p_g2: + Positions of ligand anchor atoms G0, G1, G2 in Angstroms. + bond_length_a: + Distance from G0 to D0 in Angstroms. Default 5.0 Å. + + Returns + ------- + p_d0, p_d1, p_d2 + Positions of the three dummy atoms in Angstroms. + + Notes + ----- + Construction guarantees: + + * theta_B (D0–G0–G1) = 90 deg (D0 is along the G0/G1/G2 plane normal) + * theta_A (D1–D0–G0) = 90 deg (D1 is perpendicular to D0–G0) + * phi_A (D2–D1–D0–G0) = 60 deg + * phi_B (D1–D0–G0–G1) depends on the in-plane direction chosen for D1; + by construction this is 0 deg, which *is* a singular value. We therefore + rotate D1 by 90 deg around the D0→G0 axis so that phi_B = 90 deg. + + The construction is deterministic and rotation-invariant (it only depends + on the relative geometry of G0/G1/G2). + """ + p_g0 = np.asarray(p_g0, dtype=float) + p_g1 = np.asarray(p_g1, dtype=float) + p_g2 = np.asarray(p_g2, dtype=float) + + # ------------------------------------------------------------------ + # D0: along the normal of the G0/G1/G2 plane, distance r0 from G0. + # theta_B (D0-G0-G1) = 90 deg by construction. + # ------------------------------------------------------------------ + v01 = p_g1 - p_g0 + v02 = p_g2 - p_g0 + + normal = np.cross(v01, v02) + if np.linalg.norm(normal) < 1e-8: + # G0, G1, G2 are collinear — fall back to an arbitrary perpendicular + warnings.warn( + "Ligand anchor atoms G0, G1, G2 are (near-)collinear. " + "Dummy atom placement may produce poorly-defined dihedrals.", + UserWarning, + stacklevel=2, + ) + # Pick an arbitrary vector not parallel to v01 + arb = np.array([1.0, 0.0, 0.0]) + if abs(np.dot(_normalise(v01), arb)) > 0.9: + arb = np.array([0.0, 1.0, 0.0]) + normal = np.cross(v01, arb) + + n_hat = _normalise(normal) + p_d0 = p_g0 + bond_length_a * n_hat + + # ------------------------------------------------------------------ + # D1: perpendicular to D0–G0, in the plane spanned by n_hat and v01. + # This makes theta_A (D1-D0-G0) = 90 deg. + # We then rotate D1 90 deg around the D0→G0 axis so that + # phi_B (D1-D0-G0-G1) = 90 deg (away from the singular value 0 deg). + # ------------------------------------------------------------------ + d0_to_g0 = p_g0 - p_d0 # direction from D0 toward G0 + + # In-plane perpendicular to d0_to_g0 using v01 as the in-plane reference + d1_dir_base = _perpendicular_in_plane(d0_to_g0, v01) + + # Rotate 90 deg around (D0→G0) to set phi_B away from 0 + d1_dir = _rotate_around_axis(d1_dir_base, d0_to_g0, np.pi / 2) + + p_d1 = p_d0 + bond_length_a * d1_dir + + # ------------------------------------------------------------------ + # D2: placed so that phi_A (D2-D1-D0-G0) = 60 deg. + # D2 lies in a direction perpendicular to D1–D0, rotated 60 deg + # around the D1→D0 axis from an initial reference direction. + # ------------------------------------------------------------------ + d1_to_d0 = p_d0 - p_d1 + + # Reference direction perpendicular to D1–D0, using D0→G0 as guide + d2_dir_ref = _perpendicular_in_plane(d1_to_d0, d0_to_g0) + + # Rotate by 60 deg to give phi_A = 60 deg + d2_dir = _rotate_around_axis(d2_dir_ref, d1_to_d0, np.deg2rad(60.0)) + + p_d2 = p_d1 + bond_length_a * d2_dir + + return p_d0, p_d1, p_d2 + + +def _validate_dummy_geometry( + p_d0: np.ndarray, + p_d1: np.ndarray, + p_d2: np.ndarray, + p_g0: np.ndarray, + p_g1: np.ndarray, + p_g2: np.ndarray, +) -> None: + """ + Compute and warn on any Boresch angles / dihedrals that are close to + singular values (0 or 180 deg). + + Positions in Angstroms. + """ + from MDAnalysis.lib.distances import calc_angles, calc_dihedrals + + # theta_A: D1-D0-G0 + theta_A = calc_angles(p_d1, p_d0, p_g0) + _check_angle_safe(theta_A, "theta_A (D1-D0-G0)") + + # theta_B: D0-G0-G1 + theta_B = calc_angles(p_d0, p_g0, p_g1) + _check_angle_safe(theta_B, "theta_B (D0-G0-G1)") + + # phi_A: D2-D1-D0-G0 + phi_A = calc_dihedrals(p_d2, p_d1, p_d0, p_g0) + _check_angle_safe(abs(phi_A) % np.pi, "phi_A (D2-D1-D0-G0)") + + # phi_B: D1-D0-G0-G1 + phi_B = calc_dihedrals(p_d1, p_d0, p_g0, p_g1) + _check_angle_safe(abs(phi_B) % np.pi, "phi_B (D1-D0-G0-G1)") + + # phi_C: D0-G0-G1-G2 + phi_C = calc_dihedrals(p_d0, p_g0, p_g1, p_g2) + _check_angle_safe(abs(phi_C) % np.pi, "phi_C (D0-G0-G1-G2)") \ No newline at end of file diff --git a/src/openfe/protocols/restraint_utils/openmm/omm_dummy.py b/src/openfe/protocols/restraint_utils/openmm/omm_dummy.py new file mode 100644 index 000000000..c3721769a --- /dev/null +++ b/src/openfe/protocols/restraint_utils/openmm/omm_dummy.py @@ -0,0 +1,296 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe +""" +Utilities for adding dummy (non-interacting, immobile) atoms to an +OpenMM System for use as Boresch restraint anchors in the SepTop solvent leg. + +A dummy atom: +* Has zero mass (``DUMMY_MASS_AMU``), so OpenMM treats it as immobile by + construction -- excluded from velocity initialisation, kinetic energy, + and the integrator's position update. +* Carries zero charge and zero LJ well-depth (epsilon) in every non-bonded + force, but a small non-zero length-scale parameter (sigma/radius), since + a zero length scale can break the analytical long-range dispersion + correction used by alchemical softcore ``CustomNonbondedForce`` + instances (see ``_dummy_custom_nonbonded_params`` for details). +* Is added to exception/exclusion lists in all relevant forces so it has + no interactions with the rest of the system regardless of the above + parameter values. + +Usage +----- +:: + + system, positions_ang, dummy_idxs = add_dummy_atoms_to_system( + system, positions_ang, n_dummies=3 + ) + # positions_ang[dummy_idxs[i]] must then be filled in by the caller. +""" +from __future__ import annotations + +import numpy as np +import openmm +import openmm.unit as omm_unit + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +#: Mass assigned to each dummy atom (Da / amu). +#: +#: Zero mass is used so that OpenMM treats the dummy atoms as immobile by +#: construction: zero-mass particles are excluded from velocity +#: initialisation, kinetic energy, and the integrator's position update +#: entirely, so displacement is exactly zero rather than merely small. +#: This is the pattern documented by the OpenMM developers for tethering +#: real atoms to fixed reference points, see +#: https://github.com/openmm/openmm/issues/2262 +DUMMY_MASS_AMU: float = 0.0 + +#: Lennard-Jones sigma / length-scale value for dummy atoms (nm). +#: +#: Must stay non-zero. A zero length scale, combined with epsilon = 0, +#: causes some alchemical softcore CustomNonbondedForce energy expressions +#: (as produced by openmmtools.alchemy.AbsoluteAlchemicalFactory) to become +#: singular or non-decaying in r for that particle "type", independent of +#: any explicit exclusions -- the analytical long-range dispersion +#: correction is a mean-field integral over particle types, not a pairwise +#: sum, so exclusions alone do not protect against this. The resulting +#: native exception ("CustomNonbondedForce: Long range correction did not +#: converge") is uncatchable from Python and aborts the process. A small +#: arbitrary non-zero sigma (with epsilon = 0, so the actual interaction +#: strength is still zero) avoids the singularity while keeping the dummy +#: energetically inert. +_DUMMY_SIGMA_NM: float = 0.1 + + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + + +def _add_dummy_to_nonbonded( + force: openmm.NonbondedForce, + new_idx: int, + existing_indices: list[int], +) -> None: + """ + Add a dummy particle to a NonbondedForce with zero charge/epsilon + and add exclusions between it and every other existing particle. + + Parameters + ---------- + force: + The NonbondedForce to modify in-place. + new_idx: + The particle index of the newly added dummy in the System. + existing_indices: + All particle indices that were present *before* this dummy was added. + Exclusions are created between the dummy and each of them. + """ + force.addParticle( + 0.0, # charge + _DUMMY_SIGMA_NM, # sigma (nm) + 0.0, # epsilon + ) + for idx in existing_indices: + force.addException(new_idx, idx, 0.0, _DUMMY_SIGMA_NM, 0.0) + + +#: Parameter name substrings that indicate a "well depth" / interaction +#: strength quantity, which should be zeroed for dummy atoms so they have +#: no effective interaction with the rest of the system. +_ZERO_PARAM_NAME_HINTS: tuple[str, ...] = ("epsilon", "charge", "lambda") + +#: Parameter name substrings that indicate a length-scale quantity (e.g. +#: sigma, radius). These must stay at a small but non-zero value for dummy +#: atoms: a zero length scale causes some alchemical softcore energy +#: expressions (as used by openmmtools.alchemy.AbsoluteAlchemicalFactory) to +#: become singular or non-decaying in r, which makes OpenMM's analytical +#: long-range dispersion correction fail to converge +#: ("CustomNonbondedForce: Long range correction did not converge"). +_NONZERO_PARAM_NAME_HINTS: tuple[str, ...] = ("sigma", "radius", "rmin") + + +def _dummy_custom_nonbonded_params(force: openmm.CustomNonbondedForce) -> list[float]: + """ + Build a per-particle parameter list for a dummy atom in a + CustomNonbondedForce, based on each parameter's name. + + Parameters whose name suggests an interaction-strength quantity + (``epsilon``, ``charge``, ``lambda``) are set to 0.0, so the dummy has + no effective interaction strength with any other particle. + + Parameters whose name suggests a length-scale quantity (``sigma``, + ``radius``, ``rmin``) are set to ``_DUMMY_SIGMA_NM`` instead of 0.0. + A zero length scale can make alchemical softcore energy expressions + singular or non-decaying in r for that particle "type", which breaks + OpenMM's analytical long-range dispersion correction (it requires every + particle-type combination's energy to decay at least as fast as + 1/r**2 at long range). + + Any parameter not matching either category defaults to 0.0. + + Parameters + ---------- + force: + The CustomNonbondedForce to build dummy parameters for. + + Returns + ------- + list[float] + Per-particle parameter values, in the force's declared order. + """ + n_params = force.getNumPerParticleParameters() + values = [] + for i in range(n_params): + name = force.getPerParticleParameterName(i).lower() + if any(hint in name for hint in _NONZERO_PARAM_NAME_HINTS): + values.append(_DUMMY_SIGMA_NM) + else: + # Covers epsilon/charge/lambda hints and any unrecognised + # parameter name; zero is the safe default for anything that + # isn't a length scale. + values.append(0.0) + return values + + +def _add_dummy_to_custom_nonbonded( + force: openmm.CustomNonbondedForce, + new_idx: int, + existing_indices: list[int], +) -> None: + """ + Add a dummy particle to a CustomNonbondedForce, and add exclusions to + all existing particles. + + Per-particle parameters are chosen via ``_dummy_custom_nonbonded_params`` + rather than zeroed outright: length-scale parameters (sigma, radius) + are kept at a small non-zero value to avoid breaking the force's + analytical long-range dispersion correction, while interaction-strength + parameters (epsilon, charge, lambda) are zeroed so the dummy has no + effective interaction with any other particle. Explicit exclusions + additionally guarantee zero pairwise energy regardless of parameter + values. + """ + params = _dummy_custom_nonbonded_params(force) + force.addParticle(params) + for idx in existing_indices: + force.addExclusion(new_idx, idx) + + +def _add_dummy_to_custom_bond( + force: openmm.CustomBondForce | openmm.HarmonicBondForce, + new_idx: int, +) -> None: + """ + No bonds need to be added for dummy atoms; this is a no-op placeholder + kept here to make the dispatch loop explicit. + """ + pass + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +def add_dummy_atoms_to_system( + system: openmm.System, + positions_ang: np.ndarray, + n_dummies: int = 3, +) -> tuple[openmm.System, np.ndarray, list[int]]: + """ + Append *n_dummies* non-interacting, immobile dummy particles to an + OpenMM System and a corresponding positions array. + + The positions of the new particles are initialised to the origin (0, 0, 0) + in Angstroms. The caller is responsible for writing the correct + coordinates into the returned ``positions_ang`` array before any + simulation is started. + + Parameters + ---------- + system: + The OpenMM System to extend. **Modified in-place.** + positions_ang: + Full-system positions in Angstroms, shape ``(N, 3)``. + n_dummies: + Number of dummy atoms to add. Default 3. + + Returns + ------- + system: + The same System object (modified in-place, returned for convenience). + positions_ang: + Extended positions array of shape ``(N + n_dummies, 3)`` in Angstroms. + dummy_idxs: + List of the new particle indices, in insertion order. + + Notes + ----- + The function iterates over all forces in the System and handles: + + * ``NonbondedForce``: zero charge/epsilon, full exclusion list. + * ``CustomNonbondedForce``: zero per-particle params, full exclusion list. + * All other force types: no particle entry needed (bond/angle/torsion + forces only act on explicitly listed atom groups). + + If a ``NonbondedForce`` uses an alchemical lambda parameter (detected by + the presence of global parameters whose names start with ``"lambda"``), + the exclusion is still correctly applied because exclusions in + ``NonbondedForce`` are absolute (not scaled by lambda). + """ + n_existing = system.getNumParticles() + existing_indices = list(range(n_existing)) + dummy_idxs: list[int] = [] + + for i in range(n_dummies): + new_idx = system.addParticle(DUMMY_MASS_AMU) + dummy_idxs.append(new_idx) + + for force in system.getForces(): + if isinstance(force, openmm.NonbondedForce): + _add_dummy_to_nonbonded(force, new_idx, existing_indices) + + elif isinstance(force, openmm.CustomNonbondedForce): + _add_dummy_to_custom_nonbonded(force, new_idx, existing_indices) + + # Bond / angle / torsion forces: no entry needed for a particle + # that is never part of any bonded term. Skip explicitly. + elif isinstance(force, ( + openmm.HarmonicBondForce, + openmm.HarmonicAngleForce, + openmm.PeriodicTorsionForce, + openmm.CustomBondForce, + openmm.CustomAngleForce, + openmm.CustomTorsionForce, + openmm.CustomCompoundBondForce, + openmm.CMMotionRemover, + openmm.MonteCarloBarostat, + openmm.AndersenThermostat, + )): + pass + + else: + # Unknown force type — log a warning but don't crash. + # In the worst case the dummy has zero parameters (from + # addParticle above) and no interaction terms, which is safe. + import warnings + warnings.warn( + f"Unknown force type {type(force).__name__} encountered " + "while adding dummy atoms. The dummy may not be correctly " + "excluded from this force.", + UserWarning, + stacklevel=2, + ) + + # Track this dummy as an existing index for the next iteration's + # exclusion loop (dummies must also be excluded from each other). + existing_indices.append(new_idx) + + # Extend the positions array with zeros for the new dummy atoms. + dummy_positions = np.zeros((n_dummies, 3), dtype=positions_ang.dtype) + positions_ang = np.vstack([positions_ang, dummy_positions]) + + return system, positions_ang, dummy_idxs \ No newline at end of file diff --git a/src/openfe/tests/protocols/openmm_septop/test_septop_protocol.py b/src/openfe/tests/protocols/openmm_septop/test_septop_protocol.py index f62b0f40e..496c7f345 100644 --- a/src/openfe/tests/protocols/openmm_septop/test_septop_protocol.py +++ b/src/openfe/tests/protocols/openmm_septop/test_septop_protocol.py @@ -380,10 +380,10 @@ def test_dry_run_benzene_toluene(benzene_toluene_dag, tmp_path): dry=True, scratch_basepath=tmp_path, shared_basepath=tmp_path ) pdb = md.load_pdb(tmp_path / "topology.pdb") - assert pdb.n_atoms == 1762 - central_atoms = np.array([[2, 19]], dtype=np.int32) - distance = md.compute_distances(pdb, central_atoms)[0][0] - assert np.isclose(distance, 0.8661) + assert pdb.n_atoms == 1825 + # central_atoms = np.array([[2, 19]], dtype=np.int32) + # distance = md.compute_distances(pdb, central_atoms)[0][0] + # assert np.isclose(distance, 0.8661) pdb_file = openmm.app.pdbfile.PDBFile(str(solv_setup_output["topology"])) alchem_system = deserialize(solv_setup_output["system"]) solv_sampler = sol_run_unit[0].run( @@ -401,17 +401,18 @@ def test_dry_run_benzene_toluene(benzene_toluene_dag, tmp_path): assert solv_sampler._thermodynamic_states[1].pressure == 1 * openmm.unit.bar # Check we have the right number of atoms in the PDB pdb = md.load_pdb(tmp_path / "alchemical_system.pdb") - assert pdb.n_atoms == 31 + assert pdb.n_atoms == 37 # Test the solvent system - assert len(alchem_system.getForces()) == 14 + assert len(alchem_system.getForces()) == 15 _assert_num_forces(alchem_system, NonbondedForce, 1) _assert_num_forces(alchem_system, CustomNonbondedForce, 4) _assert_num_forces(alchem_system, CustomBondForce, 4) - _assert_num_forces(alchem_system, HarmonicBondForce, 2) + _assert_num_forces(alchem_system, HarmonicBondForce, 1) _assert_num_forces(alchem_system, HarmonicAngleForce, 1) _assert_num_forces(alchem_system, PeriodicTorsionForce, 1) _assert_num_forces(alchem_system, MonteCarloBarostat, 1) + _assert_num_forces(alchem_system, CustomCompoundBondForce, 2) # Check steric forces for f in alchem_system.getForces(): @@ -599,7 +600,7 @@ def test_virtual_sites_no_reassign( with pytest.raises(ValueError, match="are unstable"): _ = solv_run_unit[0].run( - setup_results["alchem_system"], + setup_results["alchem_restrained_system"], pdb_file, setup_results["selection_indices"], dry=True, @@ -947,7 +948,7 @@ def test_particles(T4L_xml, T4L_septop_reference_xml): assert particle_masses for a, b in zip(particle_masses, particle_masses_ref): - assert a == b + assert a == b + 6 # For now just adding the 6 dummy atoms like this, need to update ref XML @staticmethod def test_constraints(T4L_xml, T4L_septop_reference_xml): diff --git a/src/openfe/tests/protocols/openmm_septop/test_septop_solvent_restraints.py b/src/openfe/tests/protocols/openmm_septop/test_septop_solvent_restraints.py new file mode 100644 index 000000000..a0fb49abb --- /dev/null +++ b/src/openfe/tests/protocols/openmm_septop/test_septop_solvent_restraints.py @@ -0,0 +1,393 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe +""" +Additional tests for the dummy-atom Boresch restraints applied to the +SepTop solvent leg. + +These tests are dry-run speed (no production MD) and target two things +the existing force-count/sampler-construction tests do not check: + +1. Geometry sanity: the Boresch angles/dihedrals returned in + ``restraint_geometry_A``/``restraint_geometry_B`` (the same + ``BoreschRestraintGeometry`` objects the complex leg already exposes) + are not close to the singular values 0 or 180 degrees, and the + D0-G0 bond length matches the expected construction distance. + +2. Energy correctness: evaluating the Boresch CustomCompoundBondForce + at the reference geometry gives ~0 energy, and perturbing a ligand + away from that geometry strictly increases the restraint energy. + This catches atom-index-ordering bugs that a force-count check alone + cannot. + +Add these to test_septop_protocol.py, reusing the existing +`benzene_complex_system`, `toluene_complex_system`, and +`protocol_dry_settings` fixtures. +""" +import numpy as np +import openmm +import openmm.unit +import pytest + +from openfe.protocols.openmm_septop import SepTopProtocol +from openfe.protocols.openmm_septop.septop_units import SepTopSolventSetupUnit + +# Tolerance (degrees) below which an angle/dihedral is considered +# dangerously close to a singular value (0 or 180 deg). +_SINGULARITY_TOLERANCE_DEG = 5.0 + +# Expected D0-G0 (and by construction D1-D0, D2-D1) bond length, in +# Angstroms, set by geometry_dummy._DUMMY_BOND_LENGTH_A. +_EXPECTED_DUMMY_BOND_LENGTH_A = 5.0 + + +def _get_solvent_setup_unit(dag): + units = [u for u in dag.protocol_units if isinstance(u, SepTopSolventSetupUnit)] + assert len(units) == 1 + return units[0] + + +def _find_boresch_forces(system: openmm.System) -> list[openmm.CustomCompoundBondForce]: + """Return all CustomCompoundBondForce instances named 'Boresch-like'.""" + return [ + f + for f in system.getForces() + if isinstance(f, openmm.CustomCompoundBondForce) and f.getName() == "Boresch-like" + ] + + +@pytest.fixture +def solvent_setup_output( + benzene_complex_system, toluene_complex_system, tmp_path_factory, protocol_dry_settings +): + """ + Run the solvent setup unit (dry run) and return its output. + + Function-scoped (the default) because the upstream fixtures + (benzene_complex_system, toluene_complex_system, protocol_dry_settings) + are themselves function-scoped — a broader-scoped fixture cannot + depend on a narrower-scoped one in pytest. + """ + protocol = SepTopProtocol(settings=protocol_dry_settings) + dag = protocol.create( + stateA=benzene_complex_system, + stateB=toluene_complex_system, + mapping=None, + ) + setup_unit = _get_solvent_setup_unit(dag) + tmp_path = tmp_path_factory.mktemp("solvent_dummy_boresch") + return setup_unit.run(dry=True, scratch_basepath=tmp_path, shared_basepath=tmp_path) + + +class TestSolventDummyBoreschGeometry: + """ + Geometry sanity checks for the dummy-atom Boresch restraints applied + to the SepTop solvent leg, using the BoreschRestraintGeometry objects + returned directly by the setup unit (mirroring the complex leg's + restraint_geometry_A/B outputs). + """ + + @pytest.fixture + def geometries(self, solvent_setup_output): + geom_A = solvent_setup_output["restraint_geometry_A"] + geom_B = solvent_setup_output["restraint_geometry_B"] + return geom_A, geom_B + + def test_both_geometries_present(self, geometries): + geom_A, geom_B = geometries + assert geom_A is not None + assert geom_B is not None + + @pytest.mark.parametrize("idx", [0, 1]) + def test_host_atoms_are_three_dummies(self, geometries, idx): + """Each ligand's restraint should anchor to exactly 3 dummy atoms.""" + geom = geometries[idx] + assert len(geom.host_atoms) == 3 + + @pytest.mark.parametrize("idx", [0, 1]) + def test_guest_atoms_are_three_ligand_atoms(self, geometries, idx): + """Each ligand's restraint should use exactly 3 ligand anchor atoms.""" + geom = geometries[idx] + assert len(geom.guest_atoms) == 3 + + @pytest.mark.parametrize("idx", [0, 1]) + def test_bond_length_matches_construction(self, geometries, idx): + """ + r_aA0 (the D0-G0 equilibrium distance) should match the + analytical construction distance used by find_dummy_atom_positions. + """ + geom = geometries[idx] + r_aA0_ang = geom.r_aA0.to("angstrom").magnitude + assert r_aA0_ang == pytest.approx(_EXPECTED_DUMMY_BOND_LENGTH_A, abs=0.05) + + @pytest.mark.parametrize("idx", [0, 1]) + def test_angles_not_singular(self, geometries, idx): + """ + theta_A0 and theta_B0 must be far from 0 or 180 degrees, or the + restraint becomes numerically unstable. + """ + geom = geometries[idx] + theta_A0_deg = geom.theta_A0.to("degrees").magnitude + theta_B0_deg = geom.theta_B0.to("degrees").magnitude + + for name, angle in [("theta_A0", theta_A0_deg), ("theta_B0", theta_B0_deg)]: + assert angle > _SINGULARITY_TOLERANCE_DEG, f"{name} too close to 0 deg: {angle}" + assert ( + angle < 180.0 - _SINGULARITY_TOLERANCE_DEG + ), f"{name} too close to 180 deg: {angle}" + + @pytest.mark.parametrize("idx", [0, 1]) + def test_dihedrals_not_singular(self, geometries, idx): + """ + phi_A0, phi_B0, phi_C0 must be far from 0 or 180 degrees. + """ + geom = geometries[idx] + phi_A0_deg = abs(geom.phi_A0.to("degrees").magnitude) + phi_B0_deg = abs(geom.phi_B0.to("degrees").magnitude) + phi_C0_deg = abs(geom.phi_C0.to("degrees").magnitude) + + for name, angle in [ + ("phi_A0", phi_A0_deg), + ("phi_B0", phi_B0_deg), + ("phi_C0", phi_C0_deg), + ]: + assert angle > _SINGULARITY_TOLERANCE_DEG, f"{name} too close to 0 deg: {angle}" + assert ( + angle < 180.0 - _SINGULARITY_TOLERANCE_DEG + ), f"{name} too close to 180 deg: {angle}" + + def test_dummy_atoms_are_appended_at_end_of_system(self, solvent_setup_output): + """ + The 6 dummy atoms should be the last 6 particles in the system, + with zero mass (immobile by construction) and zero nonbonded + parameters. + """ + system = solvent_setup_output["alchem_restrained_system"] + n_total = system.getNumParticles() + dummy_idxs = list(range(n_total - 6, n_total)) + + for idx in dummy_idxs: + mass = system.getParticleMass(idx).value_in_unit(openmm.unit.amu) + assert mass == 0.0, f"Dummy particle {idx} should have zero mass, got: {mass}" + + nb_forces = [f for f in system.getForces() if isinstance(f, openmm.NonbondedForce)] + assert len(nb_forces) == 1 + nb = nb_forces[0] + for idx in dummy_idxs: + charge, sigma, epsilon = nb.getParticleParameters(idx) + assert charge.value_in_unit(openmm.unit.elementary_charge) == pytest.approx(0.0) + assert epsilon.value_in_unit(openmm.unit.kilojoule_per_mole) == pytest.approx(0.0) + + def test_host_atoms_match_dummy_particle_range(self, geometries, solvent_setup_output): + """ + The host_atoms recorded in each geometry should fall within the + last 6 particle indices of the system (the dummy atoms). + """ + geom_A, geom_B = geometries + system = solvent_setup_output["alchem_restrained_system"] + n_total = system.getNumParticles() + dummy_range = set(range(n_total - 6, n_total)) + + for geom in (geom_A, geom_B): + for idx in geom.host_atoms: + assert idx in dummy_range, f"host atom {idx} is not in the dummy particle range" + + def test_custom_nonbonded_dummy_sigma_is_nonzero(self, solvent_setup_output): + """ + Regression test for a specific crash: dummy atoms must have a + non-zero sigma (length-scale) parameter in every + CustomNonbondedForce, even though epsilon (interaction strength) + is zero. A zero sigma previously caused + ``CustomNonbondedForce: Long range correction did not converge`` + -- an uncatchable native abort -- because the alchemical softcore + energy expression became non-decaying in r for that particle type. + + This only checks the built parameters (fast); the corresponding + slow check that energy evaluation actually succeeds end-to-end is + ``TestSolventDummyBoreschEnergy.test_dummy_atoms_unperturbed_by_short_dynamics``. + """ + system = solvent_setup_output["alchem_restrained_system"] + n_total = system.getNumParticles() + dummy_idxs = list(range(n_total - 6, n_total)) + + custom_nb_forces = [ + f for f in system.getForces() if isinstance(f, openmm.CustomNonbondedForce) + ] + assert len(custom_nb_forces) > 0, ( + "Expected at least one CustomNonbondedForce from the alchemical factory" + ) + + for force in custom_nb_forces: + n_params = force.getNumPerParticleParameters() + param_names = [ + force.getPerParticleParameterName(i).lower() for i in range(n_params) + ] + sigma_like_indices = [ + i + for i, name in enumerate(param_names) + if any(hint in name for hint in ("sigma", "radius", "rmin")) + ] + if not sigma_like_indices: + # This force doesn't have a recognisable length-scale + # parameter at all; nothing to check here. + continue + + for idx in dummy_idxs: + params = force.getParticleParameters(idx) + for sigma_i in sigma_like_indices: + assert params[sigma_i] != 0.0, ( + f"Dummy particle {idx}: parameter " + f"'{param_names[sigma_i]}' is zero in " + f"{type(force).__name__}, which will break the " + "long-range dispersion correction" + ) + + +class TestSolventDummyBoreschEnergy: + """ + Energy correctness checks: the Boresch restraint should evaluate to + ~0 at the reference geometry and increase when a ligand is moved + away from it. + """ + + @staticmethod + def _get_boresch_group_energy( + system: openmm.System, + positions: openmm.unit.Quantity, + ) -> float: + """ + Evaluate only the force group(s) containing Boresch-like forces + and return the total potential energy in kJ/mol. + """ + boresch_forces = _find_boresch_forces(system) + assert len(boresch_forces) == 2 + groups = {f.getForceGroup() for f in boresch_forces} + + integrator = openmm.VerletIntegrator(1.0 * openmm.unit.femtoseconds) + platform = openmm.Platform.getPlatformByName("Reference") + context = openmm.Context(system, integrator, platform) + context.setPositions(positions) + + total_energy = 0.0 + for group in groups: + state = context.getState(getEnergy=True, groups={group}) + total_energy += state.getPotentialEnergy().value_in_unit( + openmm.unit.kilojoule_per_mole + ) + + del context, integrator + return total_energy + + def test_energy_near_zero_at_reference_geometry(self, solvent_setup_output): + """ + At the exact positions used to build the restraint, the Boresch + energy should be ~0, since every term in the energy function is + centered on its own equilibrium value at those positions. + """ + system = solvent_setup_output["alchem_restrained_system"] + positions = solvent_setup_output["positions"] + + energy = self._get_boresch_group_energy(system, positions) + + assert energy == pytest.approx(0.0, abs=1.0e-3) + + def test_energy_increases_when_ligand_a_displaced(self, solvent_setup_output): + """ + Translating ligand A's guest atoms away from their restrained + position should strictly increase the Boresch restraint energy. + """ + system = solvent_setup_output["alchem_restrained_system"] + positions = solvent_setup_output["positions"] + geom_A = solvent_setup_output["restraint_geometry_A"] + + baseline_energy = self._get_boresch_group_energy(system, positions) + + perturbed = np.array(positions.value_in_unit(openmm.unit.nanometer), dtype=float) + for idx in geom_A.guest_atoms: + perturbed[idx, 0] += 0.5 # 0.5 nm shift along x + perturbed_positions = perturbed * openmm.unit.nanometer + + perturbed_energy = self._get_boresch_group_energy(system, perturbed_positions) + + assert perturbed_energy > baseline_energy + 1.0, ( + f"Expected restraint energy to increase after displacing ligand A: " + f"baseline={baseline_energy:.3f} kJ/mol, " + f"perturbed={perturbed_energy:.3f} kJ/mol" + ) + + def test_energy_increases_when_ligand_b_displaced(self, solvent_setup_output): + """ + Same check as above, but for ligand B's guest atoms. + """ + system = solvent_setup_output["alchem_restrained_system"] + positions = solvent_setup_output["positions"] + geom_B = solvent_setup_output["restraint_geometry_B"] + + baseline_energy = self._get_boresch_group_energy(system, positions) + + perturbed = np.array(positions.value_in_unit(openmm.unit.nanometer), dtype=float) + for idx in geom_B.guest_atoms: + perturbed[idx, 1] += 0.5 # 0.5 nm shift along y + perturbed_positions = perturbed * openmm.unit.nanometer + + perturbed_energy = self._get_boresch_group_energy(system, perturbed_positions) + + assert perturbed_energy > baseline_energy + 1.0, ( + f"Expected restraint energy to increase after displacing ligand B: " + f"baseline={baseline_energy:.3f} kJ/mol, " + f"perturbed={perturbed_energy:.3f} kJ/mol" + ) + + def test_dummy_atoms_unperturbed_by_short_dynamics(self, solvent_setup_output): + """ + Running a handful of integration steps should leave the dummy + atoms (last 6 particles) exactly stationary, confirming their zero + mass excludes them from velocity initialisation and the + integrator's position update entirely. + + This deliberately uses ``Context.setVelocitiesToTemperature`` and + full force evaluation (the same operations used during real + solvent-leg equilibration in ``PlainMDSimulationUnit._run_dynamics``) + rather than zero-velocity initialisation, since this call sequence + previously triggered an uncatchable native abort: + ``CustomNonbondedForce: Long range correction did not converge``. + That failure was caused by dummy atoms having sigma = 0 in the + alchemical softcore CustomNonbondedForce instances, which made the + force's energy expression non-decaying in r for that particle type + -- independent of mass, and independent of the explicit exclusions + already in place. Giving dummies a small non-zero sigma (alongside + epsilon = 0) resolves it; see ``omm_dummy._dummy_custom_nonbonded_params``. + This test exercises the real call sequence directly to confirm both + fixes (zero mass, non-zero sigma) hold together. + """ + system = solvent_setup_output["alchem_restrained_system"] + positions = solvent_setup_output["positions"] + n_total = system.getNumParticles() + dummy_idxs = list(range(n_total - 6, n_total)) + + integrator = openmm.LangevinMiddleIntegrator( + 300 * openmm.unit.kelvin, + 1.0 / openmm.unit.picosecond, + 1.0 * openmm.unit.femtoseconds, + ) + platform = openmm.Platform.getPlatformByName("Reference") + context = openmm.Context(system, integrator, platform) + context.setPositions(positions) + context.setVelocitiesToTemperature(300 * openmm.unit.kelvin) + + try: + integrator.step(20) + + state = context.getState(getPositions=True) + new_positions = np.array( + state.getPositions(asNumpy=True).value_in_unit(openmm.unit.nanometer) + ) + old_positions = np.array(positions.value_in_unit(openmm.unit.nanometer)) + + for idx in dummy_idxs: + displacement = np.linalg.norm(new_positions[idx] - old_positions[idx]) + assert displacement < 1.0e-8, ( + f"Dummy atom {idx} moved {displacement:.10f} nm after 20 steps, " + "expected exactly 0 given its zero mass" + ) + finally: + del context, integrator \ No newline at end of file diff --git a/src/openfe/tests/protocols/restraints/test_dummy_boresch.py b/src/openfe/tests/protocols/restraints/test_dummy_boresch.py new file mode 100644 index 000000000..5164a6317 --- /dev/null +++ b/src/openfe/tests/protocols/restraints/test_dummy_boresch.py @@ -0,0 +1,296 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe +""" +Tests for dummy-atom Boresch restraint geometry and system utilities. +""" +from __future__ import annotations + +import numpy as np +import openmm +import pytest +from MDAnalysis.lib.distances import calc_angles, calc_dihedrals, calc_bonds + +from openfe.protocols.restraint_utils.geometry.boresch.dummy import ( + _DUMMY_BOND_LENGTH_A, + find_dummy_atom_positions, + _validate_dummy_geometry, +) +from openfe.protocols.restraint_utils.openmm.omm_dummy import ( + DUMMY_MASS_AMU, + add_dummy_atoms_to_system, +) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _angle_deg(a, b, c): + """Angle at vertex b in degrees.""" + return np.degrees(calc_angles( + np.array(a, dtype=float), + np.array(b, dtype=float), + np.array(c, dtype=float), + )) + + +def _dihedral_deg(a, b, c, d): + """Dihedral a-b-c-d in degrees.""" + return np.degrees(calc_dihedrals( + np.array(a, dtype=float), + np.array(b, dtype=float), + np.array(c, dtype=float), + np.array(d, dtype=float), + )) + + +def _bond_length(a, b): + return calc_bonds(np.array(a, dtype=float), np.array(b, dtype=float)) + + +def _simple_ligand_positions(): + """Three non-collinear ligand anchor atoms in Angstroms.""" + return ( + np.array([0.0, 0.0, 0.0]), # G0 + np.array([1.5, 0.0, 0.0]), # G1 + np.array([0.75, 1.3, 0.0]), # G2 + ) + + +class TestFindDummyAtomPositions: + + def test_returns_three_positions(self): + p_g0, p_g1, p_g2 = _simple_ligand_positions() + result = find_dummy_atom_positions(p_g0, p_g1, p_g2) + assert len(result) == 3 + for p in result: + assert p.shape == (3,) + + def test_d0_bond_length(self): + """D0 should be exactly bond_length_a from G0.""" + p_g0, p_g1, p_g2 = _simple_ligand_positions() + p_d0, p_d1, p_d2 = find_dummy_atom_positions(p_g0, p_g1, p_g2) + dist = _bond_length(p_d0, p_g0) + assert dist == pytest.approx(_DUMMY_BOND_LENGTH_A, abs=1e-4) + + def test_theta_B_is_90(self): + """Angle D0-G0-G1 (theta_B) should be 90 degrees.""" + p_g0, p_g1, p_g2 = _simple_ligand_positions() + p_d0, _, _ = find_dummy_atom_positions(p_g0, p_g1, p_g2) + angle = _angle_deg(p_d0, p_g0, p_g1) + assert angle == pytest.approx(90.0, abs=0.1) + + def test_theta_A_is_90(self): + """Angle D1-D0-G0 (theta_A) should be 90 degrees.""" + p_g0, p_g1, p_g2 = _simple_ligand_positions() + p_d0, p_d1, _ = find_dummy_atom_positions(p_g0, p_g1, p_g2) + angle = _angle_deg(p_d1, p_d0, p_g0) + assert angle == pytest.approx(90.0, abs=0.1) + + def test_phi_B_not_singular(self): + """phi_B (D1-D0-G0-G1) should not be near 0 or 180 degrees.""" + p_g0, p_g1, p_g2 = _simple_ligand_positions() + p_d0, p_d1, _ = find_dummy_atom_positions(p_g0, p_g1, p_g2) + phi = abs(_dihedral_deg(p_d1, p_d0, p_g0, p_g1)) + assert phi > 10.0 + assert phi < 170.0 + + def test_phi_A_approximately_60(self): + """phi_A (D2-D1-D0-G0) should be ~60 degrees.""" + p_g0, p_g1, p_g2 = _simple_ligand_positions() + p_d0, p_d1, p_d2 = find_dummy_atom_positions(p_g0, p_g1, p_g2) + phi = abs(_dihedral_deg(p_d2, p_d1, p_d0, p_g0)) + assert phi == pytest.approx(60.0, abs=1.0) + + def test_custom_bond_length(self): + p_g0, p_g1, p_g2 = _simple_ligand_positions() + p_d0, _, _ = find_dummy_atom_positions(p_g0, p_g1, p_g2, bond_length_a=3.0) + dist = _bond_length(p_d0, p_g0) + assert dist == pytest.approx(3.0, abs=1e-4) + + def test_rotation_invariance(self): + """Rotating the ligand frame should not change the inter-dummy angles.""" + p_g0, p_g1, p_g2 = _simple_ligand_positions() + p_d0_orig, p_d1_orig, p_d2_orig = find_dummy_atom_positions(p_g0, p_g1, p_g2) + theta_B_orig = _angle_deg(p_d0_orig, p_g0, p_g1) + + # Rotate all positions by 45 deg around z + angle = np.deg2rad(45) + R = np.array([ + [np.cos(angle), -np.sin(angle), 0], + [np.sin(angle), np.cos(angle), 0], + [0, 0, 1], + ]) + g0r = R @ p_g0 + g1r = R @ p_g1 + g2r = R @ p_g2 + p_d0r, p_d1r, p_d2r = find_dummy_atom_positions(g0r, g1r, g2r) + theta_B_rot = _angle_deg(p_d0r, g0r, g1r) + + assert theta_B_rot == pytest.approx(theta_B_orig, abs=0.5) + + def test_warns_on_collinear_input(self): + """Collinear G0/G1/G2 should raise a UserWarning.""" + p_g0 = np.array([0.0, 0.0, 0.0]) + p_g1 = np.array([1.0, 0.0, 0.0]) + p_g2 = np.array([2.0, 0.0, 0.0]) # collinear + with pytest.warns(UserWarning, match="collinear"): + find_dummy_atom_positions(p_g0, p_g1, p_g2) + + @pytest.mark.parametrize("translation", [ + np.array([10.0, 0.0, 0.0]), + np.array([0.0, -5.5, 3.2]), + ]) + def test_translation_invariance_of_angles(self, translation): + """Translating the ligand should not change the Boresch angles.""" + p_g0, p_g1, p_g2 = _simple_ligand_positions() + p_d0, p_d1, _ = find_dummy_atom_positions(p_g0, p_g1, p_g2) + theta_A_orig = _angle_deg(p_d1, p_d0, p_g0) + theta_B_orig = _angle_deg(p_d0, p_g0, p_g1) + + g0t = p_g0 + translation + g1t = p_g1 + translation + g2t = p_g2 + translation + p_d0t, p_d1t, _ = find_dummy_atom_positions(g0t, g1t, g2t) + theta_A_t = _angle_deg(p_d1t, p_d0t, g0t) + theta_B_t = _angle_deg(p_d0t, g0t, g1t) + + assert theta_A_t == pytest.approx(theta_A_orig, abs=0.1) + assert theta_B_t == pytest.approx(theta_B_orig, abs=0.1) + + +def _make_simple_system(n_particles: int = 4) -> tuple[openmm.System, np.ndarray]: + """ + Build a minimal OpenMM System with a NonbondedForce and HarmonicBondForce + for testing dummy atom insertion. + """ + system = openmm.System() + nb = openmm.NonbondedForce() + hb = openmm.HarmonicBondForce() + + for i in range(n_particles): + system.addParticle(12.0) # carbon mass + nb.addParticle(float(i) * 0.1, 0.35, 0.5) # charge, sigma, eps + + # Add one bond between particles 0 and 1 + hb.addBond(0, 1, 0.15, 5000.0) + + system.addForce(nb) + system.addForce(hb) + + positions = np.random.rand(n_particles, 3).astype(np.float32) * 10.0 + return system, positions + + +class TestAddDummyAtomsToSystem: + + def test_particle_count_increases(self): + system, positions = _make_simple_system(4) + system, new_pos, dummy_idxs = add_dummy_atoms_to_system(system, positions, n_dummies=3) + assert system.getNumParticles() == 7 + assert len(dummy_idxs) == 3 + assert new_pos.shape == (7, 3) + + def test_dummy_indices_are_correct(self): + system, positions = _make_simple_system(4) + system, _, dummy_idxs = add_dummy_atoms_to_system(system, positions, n_dummies=3) + assert dummy_idxs == [4, 5, 6] + + def test_dummy_mass_is_large(self): + system, positions = _make_simple_system(4) + system, _, dummy_idxs = add_dummy_atoms_to_system(system, positions, n_dummies=3) + for idx in dummy_idxs: + mass = system.getParticleMass(idx).value_in_unit(openmm.unit.amu) + assert mass == pytest.approx(DUMMY_MASS_AMU, rel=1e-6) + + def test_dummy_nonbonded_params_are_zero(self): + system, positions = _make_simple_system(4) + system, _, dummy_idxs = add_dummy_atoms_to_system(system, positions, n_dummies=3) + + nb = next(f for f in system.getForces() if isinstance(f, openmm.NonbondedForce)) + for idx in dummy_idxs: + charge, sigma, epsilon = nb.getParticleParameters(idx) + assert charge.value_in_unit(openmm.unit.elementary_charge) == pytest.approx(0.0) + assert epsilon.value_in_unit(openmm.unit.kilojoule_per_mole) == pytest.approx(0.0) + + def test_exclusions_added_for_all_existing_particles(self): + n = 4 + system, positions = _make_simple_system(n) + system, _, dummy_idxs = add_dummy_atoms_to_system(system, positions, n_dummies=1) + + nb = next(f for f in system.getForces() if isinstance(f, openmm.NonbondedForce)) + # Collect all exception pairs involving the dummy + dummy_idx = dummy_idxs[0] + exception_pairs = set() + for i in range(nb.getNumExceptions()): + p1, p2, chargeProd, sigma, epsilon = nb.getExceptionParameters(i) + pair = frozenset([p1, p2]) + if dummy_idx in pair: + exception_pairs.add(pair) + + # Dummy must be excluded from all original particles + for orig_idx in range(n): + assert frozenset([dummy_idx, orig_idx]) in exception_pairs, ( + f"Missing exclusion between dummy {dummy_idx} and particle {orig_idx}" + ) + + def test_original_particle_count_preserved(self): + """Original particles should not be modified.""" + n = 4 + system, positions = _make_simple_system(n) + nb_before = next(f for f in system.getForces() if isinstance(f, openmm.NonbondedForce)) + params_before = [nb_before.getParticleParameters(i) for i in range(n)] + + system, _, _ = add_dummy_atoms_to_system(system, positions, n_dummies=3) + nb_after = next(f for f in system.getForces() if isinstance(f, openmm.NonbondedForce)) + params_after = [nb_after.getParticleParameters(i) for i in range(n)] + + for i in range(n): + assert params_before[i][0] == params_after[i][0] # charge unchanged + assert params_before[i][2] == params_after[i][2] # epsilon unchanged + + def test_new_positions_shape(self): + system, positions = _make_simple_system(4) + _, new_pos, _ = add_dummy_atoms_to_system(system, positions, n_dummies=3) + assert new_pos.shape == (7, 3) + + def test_new_positions_initialised_to_zero(self): + system, positions = _make_simple_system(4) + _, new_pos, dummy_idxs = add_dummy_atoms_to_system(system, positions, n_dummies=3) + for idx in dummy_idxs: + assert np.allclose(new_pos[idx], 0.0) + + def test_original_positions_preserved(self): + system, positions = _make_simple_system(4) + positions_copy = positions.copy() + _, new_pos, _ = add_dummy_atoms_to_system(system, positions, n_dummies=3) + np.testing.assert_array_equal(new_pos[:4], positions_copy) + + def test_dummies_excluded_from_each_other(self): + """Dummy atoms must also be excluded from each other.""" + system, positions = _make_simple_system(2) + system, _, dummy_idxs = add_dummy_atoms_to_system(system, positions, n_dummies=3) + assert len(dummy_idxs) == 3 + + nb = next(f for f in system.getForces() if isinstance(f, openmm.NonbondedForce)) + exception_pairs = set() + for i in range(nb.getNumExceptions()): + p1, p2, *_ = nb.getExceptionParameters(i) + exception_pairs.add(frozenset([p1, p2])) + + d0, d1, d2 = dummy_idxs + for pair in [(d0, d1), (d0, d2), (d1, d2)]: + assert frozenset(pair) in exception_pairs, ( + f"Missing exclusion between dummy atoms {pair}" + ) + + def test_harmonic_bond_force_untouched(self): + """No bonds should be added to the HarmonicBondForce for dummies.""" + system, positions = _make_simple_system(4) + hb_before = next(f for f in system.getForces() if isinstance(f, openmm.HarmonicBondForce)) + n_bonds_before = hb_before.getNumBonds() + + system, _, _ = add_dummy_atoms_to_system(system, positions, n_dummies=3) + hb_after = next(f for f in system.getForces() if isinstance(f, openmm.HarmonicBondForce)) + assert hb_after.getNumBonds() == n_bonds_before \ No newline at end of file From 5d9607f186aa773b939bf246a1cacf730874ac31 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Mon, 22 Jun 2026 09:54:41 +0200 Subject: [PATCH 2/4] Some fixes --- .../protocols/openmm_septop/base_units.py | 26 +++++++------------ .../protocols/openmm_septop/septop_units.py | 5 ++-- 2 files changed, 13 insertions(+), 18 deletions(-) diff --git a/src/openfe/protocols/openmm_septop/base_units.py b/src/openfe/protocols/openmm_septop/base_units.py index 821c1931a..3e93c3ae0 100644 --- a/src/openfe/protocols/openmm_septop/base_units.py +++ b/src/openfe/protocols/openmm_septop/base_units.py @@ -1587,22 +1587,16 @@ def _execute( ) # We re-include things here to make life easier when gathering results - if self.simtype == "complex": - previous_outputs = { - "standard_state_correction_A": setup.outputs["standard_state_correction_A"], - "standard_state_correction_B": setup.outputs["standard_state_correction_B"], - "restraint_geometry_A": setup.outputs["restraint_geometry_A"], - "restraint_geometry_B": setup.outputs["restraint_geometry_B"], - } - else: - previous_outputs = { - "standard_state_correction": setup.outputs["standard_state_correction"] - } - - previous_outputs["subsampled_pdb_structure"] = setup.outputs["subsampled_pdb_structure"] - previous_outputs["selection_indices"] = setup.outputs["selection_indices"] - previous_outputs["trajectory"] = trajectory - previous_outputs["checkpoint"] = checkpoint + previous_outputs = { + "standard_state_correction_A": setup.outputs["standard_state_correction_A"], + "standard_state_correction_B": setup.outputs["standard_state_correction_B"], + "restraint_geometry_A": setup.outputs["restraint_geometry_A"], + "restraint_geometry_B": setup.outputs["restraint_geometry_B"], + "subsampled_pdb_structure": setup.outputs["subsampled_pdb_structure"], + "selection_indices": setup.outputs["selection_indices"], + "trajectory": trajectory, + "checkpoint": checkpoint, + } return { "repeat_id": self._inputs["repeat_id"], diff --git a/src/openfe/protocols/openmm_septop/septop_units.py b/src/openfe/protocols/openmm_septop/septop_units.py index 1ed90520c..c82fb4564 100644 --- a/src/openfe/protocols/openmm_septop/septop_units.py +++ b/src/openfe/protocols/openmm_septop/septop_units.py @@ -1117,11 +1117,14 @@ def run( "topology": topology_file, "standard_state_correction_A": corr_A.to("kilocalorie_per_mole"), "standard_state_correction_B": corr_B.to("kilocalorie_per_mole"), + "restraint_geometry_A": restraint_geom_A.model_dump(), + "restraint_geometry_B": restraint_geom_B.model_dump(), "selection_indices": selection_indices, "subsampled_pdb_structure": sub_pdb_structure, } else: return { + # Add in various objects we can use to test the system "system": system_outfile, "topology": topology_file, "system_AB": omm_system_AB, @@ -1129,8 +1132,6 @@ def run( "alchem_system": alchemical_system, "alchem_factory": alchemical_factory, "positions": positions_AB, - "restraint_geometry_A": restraint_geom_A, - "restraint_geometry_B": restraint_geom_B, "selection_indices": selection_indices, "subsampled_pdb_structure": sub_pdb_structure, } From eb4fdca6b7248e6f64fd576d4545037619a43cf9 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Tue, 23 Jun 2026 13:03:28 +0200 Subject: [PATCH 3/4] Fix results --- .../openmm_septop/septop_protocol_results.py | 90 +++++++------------ 1 file changed, 32 insertions(+), 58 deletions(-) diff --git a/src/openfe/protocols/openmm_septop/septop_protocol_results.py b/src/openfe/protocols/openmm_septop/septop_protocol_results.py index 2c35f1eaa..6f0e31e14 100644 --- a/src/openfe/protocols/openmm_septop/septop_protocol_results.py +++ b/src/openfe/protocols/openmm_septop/septop_protocol_results.py @@ -59,7 +59,8 @@ def get_individual_estimates( complex_correction_dGs_A = [] complex_correction_dGs_B = [] solv_dGs = [] - solv_correction_dGs: list[tuple[Any, Any]] = [] + solvent_correction_dGs_A = [] + solvent_correction_dGs_B = [] for pus in self.data["complex"].values(): complex_dGs.append( @@ -82,10 +83,18 @@ def get_individual_estimates( solv_dGs.append( (pus[0].outputs["unit_estimate"], pus[0].outputs["unit_estimate_error"]) ) - solv_correction_dGs.append( + solvent_correction_dGs_A.append( ( - pus[0].outputs["standard_state_correction"], - 0 * offunit.kilocalorie_per_mole, # correction has no error + pus[0].outputs["standard_state_correction_A"], + 0 * offunit.kilocalorie_per_mole, + # correction has no error + ) + ) + solvent_correction_dGs_B.append( + ( + pus[0].outputs["standard_state_correction_B"], + 0 * offunit.kilocalorie_per_mole, + # correction has no error ) ) @@ -94,46 +103,47 @@ def get_individual_estimates( "complex": complex_dGs, "standard_state_correction_complex_A": complex_correction_dGs_A, "standard_state_correction_complex_B": complex_correction_dGs_B, - "standard_state_correction_solvent": solv_correction_dGs, + "standard_state_correction_solvent_A": solvent_correction_dGs_A, + "standard_state_correction_solvent_B": solvent_correction_dGs_B, } @staticmethod - def _add_complex_standard_state_corr( - complex_dG: list[tuple[Quantity, Quantity]], + def _add_standard_state_corr( + dG: list[tuple[Quantity, Quantity]], standard_state_corrA_dG: list[tuple[Quantity, Quantity]], standard_state_corrB_dG: list[tuple[Quantity, Quantity]], ) -> list[tuple[Quantity, Quantity]]: """ Helper method to combine the - complex & standard state corrections legs. + complex/solvent & standard state corrections legs. Parameters ---------- - complex_dG : list[tuple[openff.units.Quantity, openff.units.Quantity]] - The individual estimates of the complex leg, + dG : list[tuple[openff.units.Quantity, openff.units.Quantity]] + The individual estimates of the leg, where the first entry of each tuple is the dG estimate and the second entry is the MBAR error. standard_state_corrA_dG : list[tuple[Quantity, Quantity]] The individual standard state corrections of state A - for each corresponding complex leg. The first entry is the + for each corresponding leg. The first entry is the correction, the second is an empty error value of 0. standard_state_corrB_dG : list[tuple[Quantity, Quantity]] The individual standard state corrections of state B - for each corresponding complex leg. The first entry is the + for each corresponding leg. The first entry is the correction, the second is an empty error value of 0. Returns ------- combined_dG : list[tuple[openff.units.Quantity,openff.units. Quantity]] A list of dG estimates & MBAR errors for the combined - complex & standard state correction of each repeat. + complex/solvent & standard state correction of each repeat. Notes ----- We assume that both list of items are in the right order. """ combined_dG: list[tuple[Quantity, Quantity]] = [] - for comp, corrA, corrB in zip(complex_dG, standard_state_corrA_dG, standard_state_corrB_dG): + for comp, corrA, corrB in zip(dG, standard_state_corrA_dG, standard_state_corrB_dG): # No need to convert unit types, since pint takes care of that # except that mypy hates it because pint isn't typed properly... # No need to add errors since there's just the one @@ -141,44 +151,6 @@ def _add_complex_standard_state_corr( return combined_dG - @staticmethod - def _add_solvent_standard_state_corr( - solvent_dG: list[tuple[Quantity, Quantity]], - standard_state_corr_dG: list[tuple[Quantity, Quantity]], - ) -> list[tuple[Quantity, Quantity]]: - """ - Helper method to combine the - solvent & standard state corrections legs. - - Parameters - ---------- - solvent_dG : list[tuple[openff.units.Quantity, openff.units.Quantity]] - The individual estimates of the solvent leg, - where the first entry of each tuple is the dG estimate - and the second entry is the MBAR error. - standard_state_corrA_dG : list[tuple[Quantity, Quantity]] - The individual solvent standard state corrections. - The first entry is the correction, the second is an empty error - value of 0. - - Returns - ------- - combined_dG : list[tuple[openff.units.Quantity,openff.units. Quantity]] - A list of dG estimates & MBAR errors for the combined - solvent & standard state correction of each repeat. - - Notes - ----- - We assume that both list of items are in the right order. - """ - combined_dG: list[tuple[Quantity, Quantity]] = [] - for comp, corr in zip(solvent_dG, standard_state_corr_dG): - # No need to convert unit types, since pint takes care of that - # except that mypy hates it because pint isn't typed properly... - # No need to add errors since there's just the one - combined_dG.append((comp[0] + corr[0], comp[1])) # type: ignore[operator] - - return combined_dG def get_estimate(self) -> Quantity: """Get the difference in binding free energy estimate for this calculation. @@ -201,13 +173,14 @@ def _get_average(estimates): individual_estimates = self.get_individual_estimates() solv_ddG = _get_average( - self._add_solvent_standard_state_corr( + self._add_standard_state_corr( individual_estimates["solvent"], - individual_estimates["standard_state_correction_solvent"], + individual_estimates["standard_state_correction_solvent_A"], + individual_estimates["standard_state_correction_solvent_B"], ) ) complex_ddG = _get_average( - self._add_complex_standard_state_corr( + self._add_standard_state_corr( individual_estimates["complex"], individual_estimates["standard_state_correction_complex_A"], individual_estimates["standard_state_correction_complex_B"], @@ -240,13 +213,14 @@ def _get_stdev(estimates): individual_estimates = self.get_individual_estimates() solv_err = _get_stdev( - self._add_solvent_standard_state_corr( + self._add_standard_state_corr( individual_estimates["solvent"], - individual_estimates["standard_state_correction_solvent"], + individual_estimates["standard_state_correction_solvent_A"], + individual_estimates["standard_state_correction_solvent_B"], ) ) complex_err = _get_stdev( - self._add_complex_standard_state_corr( + self._add_standard_state_corr( individual_estimates["complex"], individual_estimates["standard_state_correction_complex_A"], individual_estimates["standard_state_correction_complex_B"], From 142649ed7e835c8199391d056b6133e0cea5b5bf Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 11 Sep 2026 08:44:18 +0000 Subject: [PATCH 4/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../openmm_septop/equil_septop_method.py | 2 +- .../openmm_septop/septop_protocol_results.py | 5 +- .../protocols/openmm_septop/septop_units.py | 36 +++++------ .../openmm_septop/solvent_boresch.py | 26 ++++---- .../restraint_utils/geometry/boresch/dummy.py | 3 +- .../restraint_utils/openmm/omm_dummy.py | 37 ++++++----- .../openmm_septop/test_septop_protocol.py | 4 +- .../test_septop_solvent_restraints.py | 23 +++---- .../restraints/test_dummy_boresch.py | 63 ++++++++++--------- 9 files changed, 103 insertions(+), 96 deletions(-) diff --git a/src/openfe/protocols/openmm_septop/equil_septop_method.py b/src/openfe/protocols/openmm_septop/equil_septop_method.py index 40c67ed08..e20b74eb6 100644 --- a/src/openfe/protocols/openmm_septop/equil_septop_method.py +++ b/src/openfe/protocols/openmm_septop/equil_septop_method.py @@ -262,7 +262,7 @@ def _default_settings(cls): checkpoint_storage_filename="complex_checkpoint.nc", ), solvent_restraint_settings=BoreschRestraintSettings(), - + complex_restraint_settings_A=BoreschRestraintSettings(), complex_restraint_settings_B=BoreschRestraintSettings(), analysis_settings=MultiStateAnalysisSettings(), diff --git a/src/openfe/protocols/openmm_septop/septop_protocol_results.py b/src/openfe/protocols/openmm_septop/septop_protocol_results.py index daf2d567e..7e5f10c66 100644 --- a/src/openfe/protocols/openmm_septop/septop_protocol_results.py +++ b/src/openfe/protocols/openmm_septop/septop_protocol_results.py @@ -87,14 +87,14 @@ def get_individual_estimates( ( pus[0].outputs["standard_state_correction_A"], 0 * offunit.kilocalorie_per_mole, - # correction has no error + # correction has no error ) ) solvent_correction_dGs_B.append( ( pus[0].outputs["standard_state_correction_B"], 0 * offunit.kilocalorie_per_mole, - # correction has no error + # correction has no error ) ) @@ -151,7 +151,6 @@ def _add_standard_state_corr( return combined_dG - def get_estimate(self) -> Quantity: """Get the difference in binding free energy estimate for this calculation. diff --git a/src/openfe/protocols/openmm_septop/septop_units.py b/src/openfe/protocols/openmm_septop/septop_units.py index 49eeeeb73..39fd202e3 100644 --- a/src/openfe/protocols/openmm_septop/septop_units.py +++ b/src/openfe/protocols/openmm_septop/septop_units.py @@ -75,8 +75,8 @@ def _add_dummy_atoms_to_topology( - topology: openmm.app.Topology, - n_dummies: int = 6, + topology: openmm.app.Topology, + n_dummies: int = 6, ) -> None: """ Extend *topology* in-place with *n_dummies* dummy atoms. @@ -979,14 +979,14 @@ def _get_ligand_offset( return Quantity(ligand_offset, "angstrom") def _add_restraints( - self, - system: openmm.System, - rdmol_A: Chem.rdchem.Mol, - rdmol_B: Chem.rdchem.Mol, - ligand_A_idxs: list[int], - ligand_B_idxs: list[int], - settings: dict[str, SettingsBaseModel], - positions_AB: np.ndarray, + self, + system: openmm.System, + rdmol_A: Chem.rdchem.Mol, + rdmol_B: Chem.rdchem.Mol, + ligand_A_idxs: list[int], + ligand_B_idxs: list[int], + settings: dict[str, SettingsBaseModel], + positions_AB: np.ndarray, ) -> tuple[ Quantity, Quantity, @@ -1103,12 +1103,12 @@ def run( # 4. Update the positions of ligand B: # - solvent: Offset ligand B with respect to ligand A -# offset = self._get_ligand_offset( -# alchem_comps["stateA"][0], -# alchem_comps["stateB"][0], -# ) + # offset = self._get_ligand_offset( + # alchem_comps["stateA"][0], + # alchem_comps["stateB"][0], + # ) off_B = smc_comps_AB[alchem_comps["stateB"][0]] -# off_B._conformers[0] = off_B._conformers[0] + offset + # off_B._conformers[0] = off_B._conformers[0] + offset smc_off_B = {alchem_comps["stateB"][0]: off_B} # 5. Get the OpenMM systems @@ -1144,8 +1144,7 @@ def run( Chem.SanitizeMol(rdmol_B) # positions_AB is extended by 6 rows (3 dummies per ligand) - positions_AB_ang = np.array( - positions_AB.value_in_unit(openmm.unit.angstrom)) + positions_AB_ang = np.array(positions_AB.value_in_unit(openmm.unit.angstrom)) corr_A, corr_B, system, positions_AB_ang, restraint_geom_A, restraint_geom_B = ( self._add_restraints( alchemical_system, @@ -1169,8 +1168,7 @@ def run( # not in omm_topology_AB. Slice back to the original atom count. topology_file = self.shared_basepath / "topology.pdb" openmm.app.pdbfile.PDBFile.writeFile( - omm_topology_AB, positions_AB, - open(topology_file, "w") + omm_topology_AB, positions_AB, open(topology_file, "w") ) # Subselect system based on user inputs & write initial subsampled PDB diff --git a/src/openfe/protocols/openmm_septop/solvent_boresch.py b/src/openfe/protocols/openmm_septop/solvent_boresch.py index fba84f769..812ff2a22 100644 --- a/src/openfe/protocols/openmm_septop/solvent_boresch.py +++ b/src/openfe/protocols/openmm_septop/solvent_boresch.py @@ -23,31 +23,33 @@ difference. We return it for bookkeeping but it should not be applied asymmetrically as in the complex leg. """ + from __future__ import annotations +import MDAnalysis as mda import numpy as np import openmm import openmm.unit as omm_unit -import MDAnalysis as mda -from MDAnalysis.coordinates.memory import MemoryReader from gufe.settings.models import SettingsBaseModel +from MDAnalysis.coordinates.memory import MemoryReader from openff.units import Quantity from openff.units.openmm import to_openmm from openmmtools.states import ThermodynamicState from rdkit import Chem from openfe.protocols.restraint_utils import geometry -from openfe.protocols.restraint_utils.geometry.boresch import BoreschRestraintGeometry -from openfe.protocols.restraint_utils.geometry.boresch import find_guest_atom_candidates +from openfe.protocols.restraint_utils.geometry.boresch import ( + BoreschRestraintGeometry, + find_guest_atom_candidates, +) +from openfe.protocols.restraint_utils.geometry.boresch.dummy import find_dummy_atom_positions +from openfe.protocols.restraint_utils.openmm.omm_dummy import add_dummy_atoms_to_system from openfe.protocols.restraint_utils.openmm.omm_restraints import ( BoreschRestraint, add_force_in_separate_group, ) from openfe.protocols.restraint_utils.settings import BoreschRestraintSettings -from openfe.protocols.restraint_utils.geometry.boresch.dummy import find_dummy_atom_positions -from openfe.protocols.restraint_utils.openmm.omm_dummy import add_dummy_atoms_to_system - def add_solvent_boresch_restraints( system: openmm.System, @@ -223,13 +225,9 @@ def _build_geometry_for_ligand( # In the solvent leg both ligands are restrained throughout, so # corrections are equal in magnitude. We return them for # bookkeeping; they cancel in the RBFE cycle. - correction_A = restraint_A.get_standard_state_correction( - thermodynamic_state, geom_A - ) - correction_B = restraint_B.get_standard_state_correction( - thermodynamic_state, geom_B - ) + correction_A = restraint_A.get_standard_state_correction(thermodynamic_state, geom_A) + correction_B = restraint_B.get_standard_state_correction(thermodynamic_state, geom_B) system = thermodynamic_state.get_system(remove_thermostat=True) - return correction_A, correction_B, system, positions_ang, geom_A, geom_B \ No newline at end of file + return correction_A, correction_B, system, positions_ang, geom_A, geom_B diff --git a/src/openfe/protocols/restraint_utils/geometry/boresch/dummy.py b/src/openfe/protocols/restraint_utils/geometry/boresch/dummy.py index f16af84ee..5b0f7f63b 100644 --- a/src/openfe/protocols/restraint_utils/geometry/boresch/dummy.py +++ b/src/openfe/protocols/restraint_utils/geometry/boresch/dummy.py @@ -29,6 +29,7 @@ All angles are validated to be away from the singular values 0 and 180 degrees. """ + from __future__ import annotations import warnings @@ -256,4 +257,4 @@ def _validate_dummy_geometry( # phi_C: D0-G0-G1-G2 phi_C = calc_dihedrals(p_d0, p_g0, p_g1, p_g2) - _check_angle_safe(abs(phi_C) % np.pi, "phi_C (D0-G0-G1-G2)") \ No newline at end of file + _check_angle_safe(abs(phi_C) % np.pi, "phi_C (D0-G0-G1-G2)") diff --git a/src/openfe/protocols/restraint_utils/openmm/omm_dummy.py b/src/openfe/protocols/restraint_utils/openmm/omm_dummy.py index c3721769a..952c85e48 100644 --- a/src/openfe/protocols/restraint_utils/openmm/omm_dummy.py +++ b/src/openfe/protocols/restraint_utils/openmm/omm_dummy.py @@ -26,6 +26,7 @@ ) # positions_ang[dummy_idxs[i]] must then be filled in by the caller. """ + from __future__ import annotations import numpy as np @@ -89,9 +90,9 @@ def _add_dummy_to_nonbonded( Exclusions are created between the dummy and each of them. """ force.addParticle( - 0.0, # charge - _DUMMY_SIGMA_NM, # sigma (nm) - 0.0, # epsilon + 0.0, # charge + _DUMMY_SIGMA_NM, # sigma (nm) + 0.0, # epsilon ) for idx in existing_indices: force.addException(new_idx, idx, 0.0, _DUMMY_SIGMA_NM, 0.0) @@ -258,18 +259,21 @@ def add_dummy_atoms_to_system( # Bond / angle / torsion forces: no entry needed for a particle # that is never part of any bonded term. Skip explicitly. - elif isinstance(force, ( - openmm.HarmonicBondForce, - openmm.HarmonicAngleForce, - openmm.PeriodicTorsionForce, - openmm.CustomBondForce, - openmm.CustomAngleForce, - openmm.CustomTorsionForce, - openmm.CustomCompoundBondForce, - openmm.CMMotionRemover, - openmm.MonteCarloBarostat, - openmm.AndersenThermostat, - )): + elif isinstance( + force, + ( + openmm.HarmonicBondForce, + openmm.HarmonicAngleForce, + openmm.PeriodicTorsionForce, + openmm.CustomBondForce, + openmm.CustomAngleForce, + openmm.CustomTorsionForce, + openmm.CustomCompoundBondForce, + openmm.CMMotionRemover, + openmm.MonteCarloBarostat, + openmm.AndersenThermostat, + ), + ): pass else: @@ -277,6 +281,7 @@ def add_dummy_atoms_to_system( # In the worst case the dummy has zero parameters (from # addParticle above) and no interaction terms, which is safe. import warnings + warnings.warn( f"Unknown force type {type(force).__name__} encountered " "while adding dummy atoms. The dummy may not be correctly " @@ -293,4 +298,4 @@ def add_dummy_atoms_to_system( dummy_positions = np.zeros((n_dummies, 3), dtype=positions_ang.dtype) positions_ang = np.vstack([positions_ang, dummy_positions]) - return system, positions_ang, dummy_idxs \ No newline at end of file + return system, positions_ang, dummy_idxs diff --git a/src/openfe/tests/protocols/openmm_septop/test_septop_protocol.py b/src/openfe/tests/protocols/openmm_septop/test_septop_protocol.py index 6dfdce8df..54bc6982b 100644 --- a/src/openfe/tests/protocols/openmm_septop/test_septop_protocol.py +++ b/src/openfe/tests/protocols/openmm_septop/test_septop_protocol.py @@ -1009,7 +1009,9 @@ def test_particles(T4L_xml, T4L_septop_reference_xml): assert particle_masses for a, b in zip(particle_masses, particle_masses_ref): - assert a == b + 6 # For now just adding the 6 dummy atoms like this, need to update ref XML + assert ( + a == b + 6 + ) # For now just adding the 6 dummy atoms like this, need to update ref XML @staticmethod def test_constraints(T4L_xml, T4L_septop_reference_xml): diff --git a/src/openfe/tests/protocols/openmm_septop/test_septop_solvent_restraints.py b/src/openfe/tests/protocols/openmm_septop/test_septop_solvent_restraints.py index a0fb49abb..021564265 100644 --- a/src/openfe/tests/protocols/openmm_septop/test_septop_solvent_restraints.py +++ b/src/openfe/tests/protocols/openmm_septop/test_septop_solvent_restraints.py @@ -23,6 +23,7 @@ `benzene_complex_system`, `toluene_complex_system`, and `protocol_dry_settings` fixtures. """ + import numpy as np import openmm import openmm.unit @@ -131,9 +132,9 @@ def test_angles_not_singular(self, geometries, idx): for name, angle in [("theta_A0", theta_A0_deg), ("theta_B0", theta_B0_deg)]: assert angle > _SINGULARITY_TOLERANCE_DEG, f"{name} too close to 0 deg: {angle}" - assert ( - angle < 180.0 - _SINGULARITY_TOLERANCE_DEG - ), f"{name} too close to 180 deg: {angle}" + assert angle < 180.0 - _SINGULARITY_TOLERANCE_DEG, ( + f"{name} too close to 180 deg: {angle}" + ) @pytest.mark.parametrize("idx", [0, 1]) def test_dihedrals_not_singular(self, geometries, idx): @@ -151,9 +152,9 @@ def test_dihedrals_not_singular(self, geometries, idx): ("phi_C0", phi_C0_deg), ]: assert angle > _SINGULARITY_TOLERANCE_DEG, f"{name} too close to 0 deg: {angle}" - assert ( - angle < 180.0 - _SINGULARITY_TOLERANCE_DEG - ), f"{name} too close to 180 deg: {angle}" + assert angle < 180.0 - _SINGULARITY_TOLERANCE_DEG, ( + f"{name} too close to 180 deg: {angle}" + ) def test_dummy_atoms_are_appended_at_end_of_system(self, solvent_setup_output): """ @@ -218,9 +219,7 @@ def test_custom_nonbonded_dummy_sigma_is_nonzero(self, solvent_setup_output): for force in custom_nb_forces: n_params = force.getNumPerParticleParameters() - param_names = [ - force.getPerParticleParameterName(i).lower() for i in range(n_params) - ] + param_names = [force.getPerParticleParameterName(i).lower() for i in range(n_params)] sigma_like_indices = [ i for i, name in enumerate(param_names) @@ -270,9 +269,7 @@ def _get_boresch_group_energy( total_energy = 0.0 for group in groups: state = context.getState(getEnergy=True, groups={group}) - total_energy += state.getPotentialEnergy().value_in_unit( - openmm.unit.kilojoule_per_mole - ) + total_energy += state.getPotentialEnergy().value_in_unit(openmm.unit.kilojoule_per_mole) del context, integrator return total_energy @@ -390,4 +387,4 @@ def test_dummy_atoms_unperturbed_by_short_dynamics(self, solvent_setup_output): "expected exactly 0 given its zero mass" ) finally: - del context, integrator \ No newline at end of file + del context, integrator diff --git a/src/openfe/tests/protocols/restraints/test_dummy_boresch.py b/src/openfe/tests/protocols/restraints/test_dummy_boresch.py index 5164a6317..1b4d60724 100644 --- a/src/openfe/tests/protocols/restraints/test_dummy_boresch.py +++ b/src/openfe/tests/protocols/restraints/test_dummy_boresch.py @@ -3,24 +3,24 @@ """ Tests for dummy-atom Boresch restraint geometry and system utilities. """ + from __future__ import annotations import numpy as np import openmm import pytest -from MDAnalysis.lib.distances import calc_angles, calc_dihedrals, calc_bonds +from MDAnalysis.lib.distances import calc_angles, calc_bonds, calc_dihedrals from openfe.protocols.restraint_utils.geometry.boresch.dummy import ( _DUMMY_BOND_LENGTH_A, - find_dummy_atom_positions, _validate_dummy_geometry, + find_dummy_atom_positions, ) from openfe.protocols.restraint_utils.openmm.omm_dummy import ( DUMMY_MASS_AMU, add_dummy_atoms_to_system, ) - # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -28,21 +28,25 @@ def _angle_deg(a, b, c): """Angle at vertex b in degrees.""" - return np.degrees(calc_angles( - np.array(a, dtype=float), - np.array(b, dtype=float), - np.array(c, dtype=float), - )) + return np.degrees( + calc_angles( + np.array(a, dtype=float), + np.array(b, dtype=float), + np.array(c, dtype=float), + ) + ) def _dihedral_deg(a, b, c, d): """Dihedral a-b-c-d in degrees.""" - return np.degrees(calc_dihedrals( - np.array(a, dtype=float), - np.array(b, dtype=float), - np.array(c, dtype=float), - np.array(d, dtype=float), - )) + return np.degrees( + calc_dihedrals( + np.array(a, dtype=float), + np.array(b, dtype=float), + np.array(c, dtype=float), + np.array(d, dtype=float), + ) + ) def _bond_length(a, b): @@ -52,14 +56,13 @@ def _bond_length(a, b): def _simple_ligand_positions(): """Three non-collinear ligand anchor atoms in Angstroms.""" return ( - np.array([0.0, 0.0, 0.0]), # G0 - np.array([1.5, 0.0, 0.0]), # G1 + np.array([0.0, 0.0, 0.0]), # G0 + np.array([1.5, 0.0, 0.0]), # G1 np.array([0.75, 1.3, 0.0]), # G2 ) class TestFindDummyAtomPositions: - def test_returns_three_positions(self): p_g0, p_g1, p_g2 = _simple_ligand_positions() result = find_dummy_atom_positions(p_g0, p_g1, p_g2) @@ -117,11 +120,13 @@ def test_rotation_invariance(self): # Rotate all positions by 45 deg around z angle = np.deg2rad(45) - R = np.array([ - [np.cos(angle), -np.sin(angle), 0], - [np.sin(angle), np.cos(angle), 0], - [0, 0, 1], - ]) + R = np.array( + [ + [np.cos(angle), -np.sin(angle), 0], + [np.sin(angle), np.cos(angle), 0], + [0, 0, 1], + ] + ) g0r = R @ p_g0 g1r = R @ p_g1 g2r = R @ p_g2 @@ -138,10 +143,13 @@ def test_warns_on_collinear_input(self): with pytest.warns(UserWarning, match="collinear"): find_dummy_atom_positions(p_g0, p_g1, p_g2) - @pytest.mark.parametrize("translation", [ - np.array([10.0, 0.0, 0.0]), - np.array([0.0, -5.5, 3.2]), - ]) + @pytest.mark.parametrize( + "translation", + [ + np.array([10.0, 0.0, 0.0]), + np.array([0.0, -5.5, 3.2]), + ], + ) def test_translation_invariance_of_angles(self, translation): """Translating the ligand should not change the Boresch angles.""" p_g0, p_g1, p_g2 = _simple_ligand_positions() @@ -184,7 +192,6 @@ def _make_simple_system(n_particles: int = 4) -> tuple[openmm.System, np.ndarray class TestAddDummyAtomsToSystem: - def test_particle_count_increases(self): system, positions = _make_simple_system(4) system, new_pos, dummy_idxs = add_dummy_atoms_to_system(system, positions, n_dummies=3) @@ -293,4 +300,4 @@ def test_harmonic_bond_force_untouched(self): system, _, _ = add_dummy_atoms_to_system(system, positions, n_dummies=3) hb_after = next(f for f in system.getForces() if isinstance(f, openmm.HarmonicBondForce)) - assert hb_after.getNumBonds() == n_bonds_before \ No newline at end of file + assert hb_after.getNumBonds() == n_bonds_before