Merge branch 'main' into nontransformation

news/forward_reverse_warning.rst

@@ -0,0 +1,25 @@
+**Added:**
+
+* <news item>
+
+**Changed:**
+
+* <news item>
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* Calling `get_forward_and_reverse_energy_analysis` in the RFE and AFE
+  protocols now results a warning if any results are ``None`` due to
+  MBAR convergence issues.
+
+**Security:**
+
+* <news item>

openfe/protocols/openmm_afe/equil_solvation_afe_method.py

@@ -174,13 +174,13 @@ def _get_stdev(estimates):
         # return the combined error
         return np.sqrt(vac_err**2 + solv_err**2)
 
-    def get_forward_and_reverse_energy_analysis(self) -> dict[str, list[dict[str, Union[npt.NDArray, unit.Quantity]]]]:
+    def get_forward_and_reverse_energy_analysis(self) -> dict[str, list[Optional[dict[str, Union[npt.NDArray, unit.Quantity]]]]]:
         """
         Get the reverse and forward analysis of the free energies.
 
         Returns
         -------
-        forward_reverse : dict[str, list[dict[str, Union[npt.NDArray, unit.Quantity]]]]
+        forward_reverse : dict[str, list[Optional[dict[str, Union[npt.NDArray, unit.Quantity]]]]]
             A dictionary, keyed `solvent` and `vacuum` for each leg of the
             thermodynamic cycle which each contain a list of dictionaries
             containing the forward and reverse analysis of each repeat
@@ -194,16 +194,37 @@ def get_forward_and_reverse_energy_analysis(self) -> dict[str, list[dict[str, Un
               - `forward_dDGs`, `reverse_dDGs`: unit.Quantity
                   The forward and reverse estimate uncertainty for each
                   fraction of data.
+
+            If one of the cycle leg list entries is ``None``, this indicates
+            that the analysis could not be carried out for that repeat. This
+            is most likely caused by MBAR convergence issues when attempting to
+            calculate free energies from too few samples.
+
+        Raises
+        ------
+        UserWarning
+          * If any of the forward and reverse dictionaries are ``None`` in a
+            given thermodynamic cycle leg.
         """
 
-        forward_reverse: dict[str, list[dict[str, Union[npt.NDArray, unit.Quantity]]]] = {}
+        forward_reverse: dict[str, list[Optional[dict[str, Union[npt.NDArray, unit.Quantity]]]]] = {}
 
         for key in ['solvent', 'vacuum']:
             forward_reverse[key] = [
                 pus[0].outputs['forward_and_reverse_energies']
                 for pus in self.data[key].values()
             ]
 
+            if None in forward_reverse[key]:
+                wmsg = (
+                    "One or more ``None`` entries were found in the forward "
+                    f"and reverse dictionaries of the repeats of the {key} "
+                    "calculations. This is likely caused by an MBAR convergence "
+                    "failure caused by too few independent samples when "
+                    "calculating the free energies of the 10% timeseries slice."
+                )
+                warnings.warn(wmsg)
+
         return forward_reverse
 
     def get_overlap_matrices(self) -> dict[str, list[dict[str, npt.NDArray]]]:

openfe/protocols/openmm_rfe/equil_rfe_methods.py

@@ -304,7 +304,7 @@ def get_individual_estimates(self) -> list[tuple[unit.Quantity, unit.Quantity]]:
                for pus in self.data.values()]
         return dGs
 
-    def get_forward_and_reverse_energy_analysis(self) -> list[dict[str, Union[npt.NDArray, unit.Quantity]]]:
+    def get_forward_and_reverse_energy_analysis(self) -> list[Optional[dict[str, Union[npt.NDArray, unit.Quantity]]]]:
         """
         Get a list of forward and reverse analysis of the free energies
         for each repeat using uncorrelated production samples.
@@ -317,13 +317,35 @@ def get_forward_and_reverse_energy_analysis(self) -> list[dict[str, Union[npt.ND
         The 'fractions' values are a numpy array, while the other arrays are
         Quantity arrays, with units attached.
 
+        If the list entry is ``None`` instead of a dictionary, this indicates
+        that the analysis could not be carried out for that repeat. This
+        is most likely caused by MBAR convergence issues when attempting to
+        calculate free energies from too few samples.
+
+
         Returns
         -------
-        forward_reverse : dict[str, Union[npt.NDArray, unit.Quantity]]
+        forward_reverse : list[Optional[dict[str, Union[npt.NDArray, unit.Quantity]]]]
+
+
+        Raises
+        ------
+        UserWarning
+          If any of the forward and reverse entries are ``None``.
         """
         forward_reverse = [pus[0].outputs['forward_and_reverse_energies']
                            for pus in self.data.values()]
 
+        if None in forward_reverse:
+            wmsg = (
+                "One or more ``None`` entries were found in the list of "
+                "forward and reverse analyses. This is likely caused by "
+                "an MBAR convergence failure caused by too few independent "
+                "samples when calculating the free energies of the 10% "
+                "timeseries slice."
+            )
+            warnings.warn(wmsg)
+
         return forward_reverse
 
     def get_overlap_matrices(self) -> list[dict[str, npt.NDArray]]:

openfe/tests/protocols/test_openmm_afe_slow.py

@@ -64,12 +64,15 @@ def test_openmm_run_engine(platform,
     s.solvent_output_settings.checkpoint_interval = 20 * unit.femtosecond
     s.vacuum_simulation_settings.n_replicas = 20
     s.solvent_simulation_settings.n_replicas = 20
-    s.lambda_settings.lambda_elec = \
-        [0.0, 0.25, 0.5, 0.75, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0,
-        1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
-    s.lambda_settings.lambda_vdw = \
-        [0.0, 0.0, 0.0, 0.0, 0.0, 0.05, 0.1, 0.2, 0.3, 0.4, 0.5,
-        0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95, 1.0]
+    s.lambda_settings.lambda_elec = [
+        0.0, 0.25, 0.5, 0.75, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0,
+        1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0
+    ]
+    s.lambda_settings.lambda_vdw = [
+        0.0, 0.0, 0.0, 0.0, 0.0, 0.05, 0.1, 0.2, 0.3, 0.4, 0.5,
+        0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95, 1.0
+    ]
+    s.lambda_settings.lambda_restraints = [0.0 for i in range(20)]
 
 
     protocol = openmm_afe.AbsoluteSolvationProtocol(

openfe/tests/protocols/test_openmm_afe_solvation_protocol.py

@@ -792,6 +792,19 @@ def test_get_forwards_etc(self, key, protocolresult):
             if k == 'fractions':
                 assert isinstance(far1[k], np.ndarray)
 
+    @pytest.mark.parametrize('key', ['solvent', 'vacuum'])
+    def test_get_frwd_reverse_none_return(self, key, protocolresult):
+        # fetch the first result of type key
+        data = [i for i in protocolresult.data[key].values()][0][0]
+        # set the output to None
+        data.outputs['forward_and_reverse_energies'] = None
+
+        # now fetch the analysis results and expect a warning
+        wmsg = ("were found in the forward and reverse dictionaries "
+                f"of the repeats of the {key}")
+        with pytest.warns(UserWarning, match=wmsg):
+            protocolresult.get_forward_and_reverse_energy_analysis()
+
     @pytest.mark.parametrize('key', ['solvent', 'vacuum'])
     def test_get_overlap_matrices(self, key, protocolresult):
         ovp = protocolresult.get_overlap_matrices()

openfe/tests/protocols/test_openmm_equil_rfe_protocols.py

@@ -1543,6 +1543,17 @@ def test_get_forwards_etc(self, protocolresult):
                 assert isinstance(far1[k], unit.Quantity)
                 assert far1[k].is_compatible_with(unit.kilojoule_per_mole)
 
+    def test_none_foward_reverse_energies(self, protocolresult):
+        # get the first entry's results
+        data = [i for i in protocolresult.data.values()][0][0]
+        # set the forward and reverse analysis to None
+        data.outputs['forward_and_reverse_energies'] = None
+
+        # now call the getter and expect a user warning
+        wmsg = "One or more ``None`` entries were found in"
+        with pytest.warns(UserWarning, match=wmsg):
+            protocolresult.get_forward_and_reverse_energy_analysis()
+
     def test_get_overlap_matrices(self, protocolresult):
         ovp = protocolresult.get_overlap_matrices()