From 373ab70a026c7881f665bab5f3f76e06b53a216b Mon Sep 17 00:00:00 2001 From: SimonBoothroyd Date: Thu, 16 May 2024 10:37:44 +0200 Subject: [PATCH] [SepTop] Support PMX-style receptor reference atom selection --- examples/eralpha/config-septop.yaml | 3 + femto/fe/config.py | 5 + femto/fe/reference.py | 224 ++++++++++++++++++++++++++-- femto/fe/septop/_config.py | 11 ++ femto/fe/septop/_setup.py | 69 ++++++--- femto/fe/tests/test_reference.py | 16 +- femto/md/restraints.py | 16 +- 7 files changed, 304 insertions(+), 40 deletions(-) diff --git a/examples/eralpha/config-septop.yaml b/examples/eralpha/config-septop.yaml index a5c460b..5cc3a42 100644 --- a/examples/eralpha/config-septop.yaml +++ b/examples/eralpha/config-septop.yaml @@ -25,6 +25,9 @@ complex: k_dihedral_c: 20.0 kcal * rad**-2 * mol**-1 scale_k_angle_a: true + atom_selection_receptor: baumann + atom_selection_ligand: baumann + apply_hmr: true hydrogen_mass: 1.5 Da diff --git a/femto/fe/config.py b/femto/fe/config.py index 97286d8..273d9ec 100644 --- a/femto/fe/config.py +++ b/femto/fe/config.py @@ -19,6 +19,11 @@ restraints.""" +ReceptorReferenceMethod = typing.Literal["baumann", "pmx"] +"""The method to use when automatically selecting protein atoms to use in ligand +alignment restraints.""" + + class FEP(femto.md.utils.models.BaseModel): """Configure modifying a system to be scalable by FEP lambdas.""" diff --git a/femto/fe/reference.py b/femto/fe/reference.py index 2adb0c0..319d899 100644 --- a/femto/fe/reference.py +++ b/femto/fe/reference.py @@ -36,6 +36,13 @@ _RMSF_CUTOFF = 0.1 # nm +# taken from PMX at commit 3b135ad +_PMX_ANGLE_CHECK_T = 298.15 * openmm.unit.kelvin +_PMX_ANGLE_CHECK_RT = ( + openmm.unit.MOLAR_GAS_CONSTANT_R * _PMX_ANGLE_CHECK_T +).in_units_of(openmm.unit.kilojoules_per_mole) +_PMX_ANGLE_CHECK_THRESHOLD = 5.0 + def _is_angle_linear(coords: numpy.ndarray, idxs: tuple[int, int, int]) -> bool: """Check if angle is within 10 kT from 0 or 180 following the SepTop reference @@ -611,9 +618,30 @@ def _structure_to_mdtraj(structure: parmed.Structure) -> mdtraj.Trajectory: return trajectory -def select_receptor_idxs( - receptor: parmed.Structure | mdtraj.Trajectory, - ligand: parmed.Structure | mdtraj.Trajectory, +def _check_pmx_angle(angle: float) -> bool: + """Ensure that an angle is not close to 0 or 180 degrees based on the PMX restraint + selection criteria. + + Args: + angle: The angle [rad] to check. + + Returns: + Whether the angle is valid. + """ + # taken from PMX at commit 3b135ad + k_angle = 41.84 * openmm.unit.kilojoules_per_mole + + check_0 = 0.5 * k_angle * (angle - 0.0) ** 2 / _PMX_ANGLE_CHECK_RT + check_1 = 0.5 * k_angle * (angle - numpy.pi) ** 2 / _PMX_ANGLE_CHECK_RT + + return ( + check_0 >= _PMX_ANGLE_CHECK_THRESHOLD and check_1 >= _PMX_ANGLE_CHECK_THRESHOLD + ) + + +def select_receptor_idxs_baumann( + receptor: mdtraj.Trajectory, + ligand: mdtraj.Trajectory, ligand_ref_idxs: tuple[int, int, int], ) -> tuple[int, int, int]: """Select possible protein atoms for Boresch-style restraints using the method @@ -626,21 +654,12 @@ def select_receptor_idxs( Args: receptor: The receptor structure. ligand: The ligand structure. - ligand_ref_idxs: The indices of the three ligands atoms that will be restrained. + ligand_ref_idxs: The indices of the three ligand atoms that will be restrained. Returns: - The indices of the three atoms to use for the restraint + The indices of the three atoms to use for the restraint, labeled P1, P2, and P3 + respectively in the publication. """ - if not (isinstance(receptor, type(ligand)) or isinstance(ligand, type(receptor))): - raise ValueError("receptor and ligand must be the same type") - - if isinstance(receptor, parmed.Structure) and isinstance(ligand, parmed.Structure): - receptor = _structure_to_mdtraj(receptor) - ligand = _structure_to_mdtraj(ligand) - - assert ( - receptor.n_frames == ligand.n_frames - ), "receptor and ligand must have the same number of frames" receptor_idxs = _filter_receptor_atoms(receptor, ligand, ligand_ref_idxs[0]) @@ -700,6 +719,181 @@ def select_receptor_idxs( return found_r1, found_r2, found_r3 +def _select_pmx_receptor_atom( + xyz_ref: numpy.ndarray, + receptor: parmed.Structure | mdtraj.Trajectory, + receptor_ref_idxs: list[int], + ligand: parmed.Structure | mdtraj.Trajectory, + ligand_ref_idxs: tuple[int, int, int], + angle_check_type: int = 0, +) -> int: + """Select the receptor atom closest to a reference coordinate, optionally checking + angle criteria. + + Args: + xyz_ref: The reference coordinates [nm]. + angle_check_type: The type of angle to check: 0 for none, 1 for L2-L1-R angle, + 2 for L1-R1-R angle. + + Returns: + The index of the selected receptor atom. + """ + assert angle_check_type in {0, 1, 2}, "angle_check_type must be 0, 1, or 2" + + min_distance = numpy.inf + min_receptor_idx = None + + receptor_idxs = receptor.top.select("protein and (backbone or name CA or name CB)") + + def _compute_mean_angle( + xyz_a: numpy.ndarray, xyz_b: numpy.ndarray, xyz_c: numpy.ndarray + ): + angle = femto.md.utils.geometry.compute_angles( + numpy.stack([xyz_a, xyz_b, xyz_c], axis=1), + numpy.array([[0, 1, 2]]), + ).mean() + + return angle + + for receptor_idx in receptor_idxs: + if receptor_idx in receptor_ref_idxs: + continue + + xyz_receptor = receptor.xyz[:, receptor_idx, :] + + distance_sqr = ((xyz_ref - xyz_receptor) ** 2).sum(axis=-1) + distance = numpy.sqrt(distance_sqr).mean(axis=0) + # TODO: allow an apo structure to be provided and compare the distance to the + # closest atom in the apo structure. + + is_valid_angle = True + + if angle_check_type == 1: + # check l2-l1-r angle + angle = _compute_mean_angle( + ligand.xyz[:, ligand_ref_idxs[1], :], + ligand.xyz[:, ligand_ref_idxs[0], :], + receptor.xyz[:, receptor_idx, :], + ) + is_valid_angle = _check_pmx_angle(angle) + elif angle_check_type == 2: + # check l1-r1-r angle + angle = _compute_mean_angle( + ligand.xyz[:, ligand_ref_idxs[0], :], + receptor.xyz[:, receptor_ref_idxs[0], :], + receptor.xyz[:, receptor_idx, :], + ) + is_valid_angle = _check_pmx_angle(angle) + + if distance < min_distance and is_valid_angle: + min_distance = distance + min_receptor_idx = receptor_idx + + if min_receptor_idx is None: + raise ValueError("no suitable receptor atoms could be found") + + return int(min_receptor_idx) + + +def select_receptor_idxs_pmx( + receptor: mdtraj.Trajectory, + ligand: mdtraj.Trajectory, + ligand_ref_idxs: tuple[int, int, int], +) -> tuple[int, int, int]: + """Select possible protein atoms for Boresch-style restraints using the method + defined in the PMX repository. + + Notes: + The implementation is based on commit 3b135ad. + + Args: + receptor: The receptor structure. + ligand: The ligand structure. + ligand_ref_idxs: The indices of the three ligand atoms that will be restrained. + + Returns: + The indices of the three atoms to use for the restraint, where the distance + restraint should be placed between ``receptor_ref_idxs[0]`` and + ``ligand_ref_idxs[0]``. + """ + receptor_ref_idxs = [] + + r1 = _select_pmx_receptor_atom( + ligand.xyz[:, ligand_ref_idxs[0], :], + receptor, + receptor_ref_idxs, + ligand, + ligand_ref_idxs, + angle_check_type=1, + ) + receptor_ref_idxs.append(r1) + + r2 = _select_pmx_receptor_atom( + receptor.xyz[:, r1, :], + receptor, + receptor_ref_idxs, + ligand, + ligand_ref_idxs, + angle_check_type=0, + ) + receptor_ref_idxs.append(r2) + + r3 = _select_pmx_receptor_atom( + receptor.xyz[:, r2, :], + receptor, + receptor_ref_idxs, + ligand, + ligand_ref_idxs, + angle_check_type=2, + ) + receptor_ref_idxs.append(r3) + + return r1, r2, r3 + + +def select_receptor_idxs( + receptor: parmed.Structure | mdtraj.Trajectory, + ligand: parmed.Structure | mdtraj.Trajectory, + ligand_ref_idxs: tuple[int, int, int], + method: femto.fe.config.ReceptorReferenceMethod = "baumann", +) -> tuple[int, int, int]: + """Select possible protein atoms for Boresch-style restraints. + + Notes: + See ``select_receptor_idxs_baumann`` and ``select_receptor_idxs_pmx`` for more + details. + + Args: + receptor: The receptor structure. + ligand: The ligand structure. + ligand_ref_idxs: The indices of the three ligand atoms that will be restrained. + method: The method to use to select the reference atoms. + + Returns: + The indices of the three atoms to use for the restraint. The ordering is such + that where the distance restraint should be placed between + ``receptor_ref_idxs[0]`` and ``ligand_ref_idxs[0]``. + """ + + if not (isinstance(receptor, type(ligand)) or isinstance(ligand, type(receptor))): + raise ValueError("receptor and ligand must be the same type") + + if isinstance(receptor, parmed.Structure) and isinstance(ligand, parmed.Structure): + receptor = _structure_to_mdtraj(receptor) + ligand = _structure_to_mdtraj(ligand) + + assert ( + receptor.n_frames == ligand.n_frames + ), "receptor and ligand must have the same number of frames" + + if method.lower() == "pmx": + return select_receptor_idxs_pmx(receptor, ligand, ligand_ref_idxs) + elif method.lower() == "baumann": + return select_receptor_idxs_baumann(receptor, ligand, ligand_ref_idxs) + + raise NotImplementedError(f"unknown method: {method}") + + def check_receptor_idxs( receptor: parmed.Structure | mdtraj.Trajectory, receptor_idxs: tuple[int, int, int], diff --git a/femto/fe/septop/_config.py b/femto/fe/septop/_config.py index a5a3514..61852a6 100644 --- a/femto/fe/septop/_config.py +++ b/femto/fe/septop/_config.py @@ -76,6 +76,17 @@ class SepTopComplexRestraints(femto.md.config.BoreschRestraint): "based upon the *initial* distance between r3 and l1.", ) + atom_selection_receptor: femto.fe.config.ReceptorReferenceMethod = pydantic.Field( + "baumann", + description="The method to use when automatically selecting the protein atoms " + "to use in the restraint.", + ) + atom_selection_ligand: typing.Literal["baumann"] = pydantic.Field( + "baumann", + description="The method to use when automatically selecting the ligand atoms " + "to use in the restraint.", + ) + class SepTopSolutionRestraints(BaseModel): """Configure the restraints to apply in the solution phase.""" diff --git a/femto/fe/septop/_setup.py b/femto/fe/septop/_setup.py index be22e01..bdb5989 100644 --- a/femto/fe/septop/_setup.py +++ b/femto/fe/septop/_setup.py @@ -201,6 +201,44 @@ def _setup_system( return system, topology, ligand_1_ref_idxs, ligand_2_ref_idxs +def _select_receptor_ref_atoms( + receptor: parmed.Structure, + receptor_ref_query: tuple[str, str, str], + ligand_1: parmed.Structure, + ligand_1_ref_idxs: tuple[int, int, int], + ligand_2: parmed.Structure | None, + ligand_2_ref_idxs: tuple[int, int, int] | None, + restraint_config: "femto.fe.septop.SepTopComplexRestraints", +) -> tuple[tuple[int, int, int], tuple[int, int, int] | None]: + """Selects the receptor atoms to use in the alignment restraint.""" + + if receptor_ref_query is not None: + ref_idxs = femto.fe.reference.queries_to_idxs(receptor, receptor_ref_query) + return ref_idxs, None if ligand_2 is None else ref_idxs + + method = restraint_config.atom_selection_receptor.lower() + _LOGGER.info(f"selecting receptor reference atoms with method={method}") + + ref_idxs_1 = femto.fe.reference.select_receptor_idxs( + receptor, ligand_1, ligand_1_ref_idxs, method + ) + ref_idxs_2 = ref_idxs_1 + + if method != "baumann" or ligand_2 is None: + return ref_idxs_1, None if ligand_2 is None else ref_idxs_1 + + if femto.fe.reference.check_receptor_idxs( + receptor, ref_idxs_1, ligand_2, ligand_1_ref_idxs + ): + _LOGGER.info("selecting alternate receptor reference atoms for ligand 2") + + ref_idxs_2 = femto.fe.reference.select_receptor_idxs( + receptor, ligand_2, ligand_2_ref_idxs + ) + + return ref_idxs_1, ref_idxs_2 + + def setup_complex( config: "femto.fe.septop.SepTopSetupStage", receptor: parmed.amber.AmberParm, @@ -232,28 +270,15 @@ def setup_complex( idx_offset = sum(len(ligand.atoms) for ligand in ligands) _LOGGER.info("applying restraints.") - - if receptor_ref_query is None: - _LOGGER.info("selecting receptor reference atoms") - - receptor_ref_idxs_1 = femto.fe.reference.select_receptor_idxs( - receptor, ligand_1, ligand_1_ref_idxs - ) - receptor_ref_idxs_2 = receptor_ref_idxs_1 - - if ligand_2 is not None and not femto.fe.reference.check_receptor_idxs( - receptor, receptor_ref_idxs_1, ligand_2, ligand_1_ref_idxs - ): - _LOGGER.info("selecting alternate receptor reference atoms for ligand 2") - receptor_ref_idxs_2 = femto.fe.reference.select_receptor_idxs( - receptor, ligand_2, ligand_2_ref_idxs - ) - - else: - receptor_ref_idxs_1 = femto.fe.reference.queries_to_idxs( - receptor, receptor_ref_query - ) - receptor_ref_idxs_2 = receptor_ref_idxs_1 + receptor_ref_idxs_1, receptor_ref_idxs_2 = _select_receptor_ref_atoms( + receptor, + receptor_ref_query, + ligand_1, + ligand_1_ref_idxs, + ligand_2, + ligand_2_ref_idxs, + restraint_config, + ) _LOGGER.info(f"receptor ref idxs for ligand 1={receptor_ref_idxs_1}") receptor_ref_idxs_1 = tuple(i + idx_offset for i in receptor_ref_idxs_1) diff --git a/femto/fe/tests/test_reference.py b/femto/fe/tests/test_reference.py index 7a20361..61a85fe 100644 --- a/femto/fe/tests/test_reference.py +++ b/femto/fe/tests/test_reference.py @@ -524,12 +524,14 @@ def test_is_valid_r3( assert spied_trans.spy_return == expected_dihedral_trans -def test_select_receptor_idxs(cdk2_receptor, cdk2_ligand_1, cdk2_ligand_1_ref_idxs): +def test_select_receptor_idxs_baumann( + cdk2_receptor, cdk2_ligand_1, cdk2_ligand_1_ref_idxs +): # computed using the reference SepTop implementation at commit 3705ba5 expected_receptor_idxs = 830, 841, 399 receptor_idxs = femto.fe.reference.select_receptor_idxs( - cdk2_receptor, cdk2_ligand_1, cdk2_ligand_1_ref_idxs + cdk2_receptor, cdk2_ligand_1, cdk2_ligand_1_ref_idxs, method="baumann" ) assert receptor_idxs == expected_receptor_idxs @@ -538,6 +540,16 @@ def test_select_receptor_idxs(cdk2_receptor, cdk2_ligand_1, cdk2_ligand_1_ref_id ) +def test_select_receptor_idxs_pmx(cdk2_receptor, cdk2_ligand_1, cdk2_ligand_1_ref_idxs): + # visually inspected to ensure the selection looks sensible. + expected_receptor_idxs = 250, 247, 251 + + receptor_idxs = femto.fe.reference.select_receptor_idxs( + cdk2_receptor, cdk2_ligand_1, cdk2_ligand_1_ref_idxs, method="pmx" + ) + assert receptor_idxs == expected_receptor_idxs + + def test_select_protein_cavity_atoms(cdk2_receptor, cdk2_ligand_1, cdk2_ligand_2): ligands = [cdk2_ligand_1, cdk2_ligand_2] diff --git a/femto/md/restraints.py b/femto/md/restraints.py index 055ef12..769f709 100644 --- a/femto/md/restraints.py +++ b/femto/md/restraints.py @@ -183,7 +183,21 @@ def create_boresch_restraint( coords: openmm.unit.Quantity, ctx_parameter: str | None = None, ) -> openmm.CustomCompoundBondForce: - """Creates a Boresch restraint force useful in aligning a receptor and ligand. + """Creates a Boresch-style restraint force useful in aligning a receptor and ligand. + + This includes: + + * a distance restraint between ``receptor_atoms[2]`` and ``ligand_atoms[0]``, + * an angle restraint between ``receptor_atoms[1]``, ``receptor_atoms[2]``, and + ``ligand_atoms[0]``, + * an angle restraint between ``receptor_atoms[2]``, ``ligand_atoms[0]``, and + ``ligand_atoms[1]``, + * a dihedral restraint between ``receptor_atoms[0]``, ``receptor_atoms[1]``, + ``receptor_atoms[2]``, and ``ligand_atoms[0]``, + * a dihedral restraint between ``receptor_atoms[1]``, ``receptor_atoms[2]``, + ``ligand_atoms[0]``, and ``ligand_atoms[1]``, + * a dihedral restraint between ``receptor_atoms[2]``, ``ligand_atoms[0]``, + ``ligand_atoms[1]``, and ``ligand_atoms[2]``. Args: config: The restraint configuration.