Merge pull request #360 from fjclark/feature-robust-restraint-search

Improve robustness of restraint search and fix torsional force constant bug. [ci skip]

python/BioSimSpace/Sandpit/Exscientia/FreeEnergy/_restraint_search.py

@@ -590,6 +590,14 @@ def analyse(
                 f"traj {type(traj)} must be of type 'BioSimSpace.Trajectory._trajectory.Trajectory'"
             )
 
+        n_frames = traj.nFrames()
+        if not n_frames >= 50:
+            _warnings.warn(
+                f"The trajectory for restraint selection has less than 50 frames ({n_frames} frames). "
+                "This may result in poor restraint selection with excessively high force constants. "
+                "Consider running a longer simulation or saving more frames."
+            )
+
         if not isinstance(temperature, _Temperature):
             raise ValueError(
                 f"temperature {type(temperature)} must be of type 'BioSimSpace.Types.Temperature'"
@@ -667,6 +675,18 @@ def analyse(
         decoupled_mol = _system.getDecoupledMolecules()[0]
         decoupled_resname = decoupled_mol.getResidues()[0].name()
 
+        # Warn the user if the decoupled molecule is water or a macromoelcule
+        if decoupled_mol.isWater():
+            _warnings.warn(
+                "The decoupled molecule is water. Ensure that this is intended. Using constrained water hydrogens as anchor "
+                "points may produce instabilities."
+            )
+
+        if decoupled_mol.nResidues() > 1:
+            _warnings.warn(
+                "The decoupled molecule has more than one residue and is likely a macromolecule. Ensure that this is intended."
+            )
+
         ligand_selection_str = f"((resname {decoupled_resname}) and (not name H*))"
         if append_to_ligand_selection:
             ligand_selection_str += " and "
@@ -1012,6 +1032,20 @@ def _findOrderedPairs(u, ligand_selection_str, receptor_selection_str, cutoff):
         """
 
         lig_selection = u.select_atoms(ligand_selection_str)
+
+        # Ensure that there are no shared atoms between the ligand selection and all possible receptor selections.
+        all_possible_receptor_selection = u.select_atoms(receptor_selection_str)
+        shared_atoms = [
+            atom for atom in lig_selection if atom in all_possible_receptor_selection
+        ]
+        if shared_atoms:
+            raise _AnalysisError(
+                "Shared atoms between ligand and receptor selections detected. "
+                "Please ensure that you are decoupling the intended molecule and that "
+                "the ligand and receptor selections are mutually exclusive. "
+                f"Shared atoms: {shared_atoms}"
+            )
+
         pair_variance_dict = {}
 
         # Get all receptor atoms within specified distance of cutoff
@@ -1488,7 +1522,9 @@ def _findOrderedBoresch(
                             dtheta = abs(val - circmean)
                             corrected_values.append(min(dtheta, 2 * _np.pi - dtheta))
                         corrected_values = _np.array(corrected_values)
-                        boresch_dof_data[pair][dof]["var"] = corrected_values.var()
+                        boresch_dof_data[pair][dof]["var"] = _np.mean(
+                            corrected_values**2
+                        )
 
                     # Assume Gaussian distributions and calculate force constants for harmonic potentials
                     # so as to reproduce these distributions at 298 K

tests/Sandpit/Exscientia/FreeEnergy/test_restraint_search.py

@@ -2,6 +2,8 @@
 
 import numpy as np
 
+from functools import partial
+
 import BioSimSpace.Sandpit.Exscientia as BSS
 from BioSimSpace.Sandpit.Exscientia.Align import decouple
 from BioSimSpace.Sandpit.Exscientia._Exceptions import AnalysisError
@@ -143,18 +145,20 @@ class TestBSS_analysis:
     """Test selection of restraints using the inbuilt BSS method."""
 
     @staticmethod
-    @pytest.fixture(scope="class")
-    def _restraint_search(tmp_path_factory):
+    def _restraint_search_general(tmp_path_factory, ligand_idx=1):
         outdir = tmp_path_factory.mktemp("out")
         system = BSS.IO.readMolecules(
             [
                 f"{url}/crd.gro.bz2",
                 f"{url}/complex.top.bz2",
             ]
         )
-        ligand = system.getMolecule(1)
-        decoupled_ligand = decouple(ligand)
         protein = system.getMolecule(0)
+        if ligand_idx == 0:
+            ligand = protein.copy()
+        else:
+            ligand = system.getMolecule(ligand_idx)
+        decoupled_ligand = decouple(ligand)
         new_system = (protein + decoupled_ligand).toSystem()
 
         protocol = BSS.Protocol.Production()
@@ -167,6 +171,15 @@ def _restraint_search(tmp_path_factory):
         )
         return restraint_search, outdir
 
+    @pytest.fixture(scope="class")
+    def _restraint_search(self, tmp_path_factory):
+        return partial(self._restraint_search_general, tmp_path_factory, ligand_idx=1)()
+
+    @pytest.fixture(scope="class")
+    def _restraint_search_protein(self, tmp_path_factory):
+        # Select the protein as the ligand.
+        return partial(self._restraint_search_general, tmp_path_factory, ligand_idx=0)()
+
     @staticmethod
     @pytest.fixture(scope="class")
     def boresch_restraint(_restraint_search):
@@ -210,48 +223,48 @@ def test_plots_boresch(self, boresch_restraint):
     def test_dG_off_boresch(self, boresch_restraint):
         """Test if the restraint generated has the same energy"""
         restraint, _ = boresch_restraint
-        assert np.isclose(-9.8955, restraint.correction.value(), atol=0.01)
+        assert np.isclose(-9.2133, restraint.correction.value(), atol=0.01)
 
     def test_bond_boresch(self, boresch_restraint):
         restraint, _ = boresch_restraint
         equilibrium_values_r0 = (
             restraint._restraint_dict["equilibrium_values"]["r0"] / nanometer
         )
-        assert np.isclose(0.6057, equilibrium_values_r0, atol=0.001)
+        assert np.isclose(0.4987, equilibrium_values_r0, atol=0.001)
 
     def test_angles_boresch(self, boresch_restraint):
         restraint, _ = boresch_restraint
         equilibrium_values_thetaA0 = (
             restraint._restraint_dict["equilibrium_values"]["thetaA0"] / degree
         )
-        assert np.isclose(140.6085, equilibrium_values_thetaA0, atol=0.001)
+        assert np.isclose(119.3375, equilibrium_values_thetaA0, atol=0.001)
         equilibrium_values_thetaB0 = (
             restraint._restraint_dict["equilibrium_values"]["thetaB0"] / degree
         )
-        assert np.isclose(56.4496, equilibrium_values_thetaB0, atol=0.001)
+        assert np.isclose(100.8382, equilibrium_values_thetaB0, atol=0.001)
 
     def test_dihedrals_boresch(self, boresch_restraint):
         restraint, _ = boresch_restraint
         equilibrium_values_phiA0 = (
             restraint._restraint_dict["equilibrium_values"]["phiA0"] / degree
         )
-        assert np.isclose(21.6173, equilibrium_values_phiA0, atol=0.001)
+        assert np.isclose(29.1947, equilibrium_values_phiA0, atol=0.001)
         equilibrium_values_phiB0 = (
             restraint._restraint_dict["equilibrium_values"]["phiB0"] / degree
         )
-        assert np.isclose(-19.4394, equilibrium_values_phiB0, atol=0.001)
+        assert np.isclose(-136.9009, equilibrium_values_phiB0, atol=0.001)
         equilibrium_values_phiC0 = (
             restraint._restraint_dict["equilibrium_values"]["phiC0"] / degree
         )
-        assert np.isclose(71.3148, equilibrium_values_phiC0, atol=0.001)
+        assert np.isclose(-102.3908, equilibrium_values_phiC0, atol=0.001)
 
     def test_index_boresch(self, boresch_restraint):
         restraint, _ = boresch_restraint
         idxs = {
             k: restraint._restraint_dict["anchor_points"][k].index()
             for k in restraint._restraint_dict["anchor_points"]
         }
-        assert idxs == {"r1": 1560, "r2": 1558, "r3": 1562, "l1": 10, "l2": 9, "l3": 11}
+        assert idxs == {"r1": 1560, "r2": 1558, "r3": 1562, "l1": 8, "l2": 9, "l3": 13}
 
     def test_force_constant_boresch(self, boresch_restraint_k20):
         restraint, _ = boresch_restraint_k20
@@ -268,6 +281,22 @@ def test_analysis_failure_boresch(self, _restraint_search):
                 cutoff=0.1 * angstrom,
             )
 
+    def test_non_overlapping_anchor_selections_forbidden(
+        self, _restraint_search_protein
+    ):
+        """
+        Ensure that anchor atoms cannot be in the same molecule
+        (the protein has been selected as the ligand, so we are
+        guaranteed that anchor atoms will be in the same molecule).
+        """
+        restraint_search, _ = _restraint_search_protein
+        with pytest.raises(AnalysisError):
+            restraint = restraint_search.analyse(
+                method="BSS",
+                restraint_type="Boresch",
+                block=False,
+            )
+
     def test_plots_mdr(self, multiple_distance_restraint):
         """Test if all the plots have been generated correctly"""
         restraint, outdir = multiple_distance_restraint