Move methods in cinnabar.stats concerning MLE to the MLEEstimator class

cinnabar/estimators.py

@@ -17,11 +17,19 @@
 from openff.units import Quantity
 
 from cinnabar import stats
+from cinnabar._due import Doi, due
 from cinnabar.measurements import Measurement, ReferenceState
 
 if TYPE_CHECKING:
     from cinnabar.femap import FEMap  # pragma: no cover
 
+due.cite(
+    Doi("10.1021/acs.jcim.9b00528"),
+    description="Compute maximum likelihood estimate of free energies and covariance in their estimates",
+    path="cinnabar.estimators.MLEEstimator.mle",
+    cite_module=True,
+)
+
 
 @dataclass
 class EstimatorResult:
@@ -233,23 +241,22 @@ def _estimate(
             Contains :attr:`~MLEEstimatorResult.covariance_matrix` and
             :attr:`~MLEEstimatorResult.ligand_order`.
         """
-        # TODO: replace stats.mle call with a self-contained implementation
-        g, u = _build_graph_from_measurements(measurements)
+        graph, unit = self._build_graph_from_measurements(measurements)
 
-        f_i_calc, C_calc = stats.mle(g, factor="calc_DDG")
+        f_i_calc, C_calc = self.mle(graph, edge_data_label="calc_DDG")
         variance = np.diagonal(C_calc) ** 0.5
 
         ref = ReferenceState(label=source)
-        ligand_order = list(g.nodes)
+        ligand_order = list(graph.nodes)
 
         out_measurements: list[Measurement] = []
-        for n, f_i, df_i in zip(ligand_order, f_i_calc, variance):
+        for label, f_i, df_i in zip(ligand_order, f_i_calc, variance):
             out_measurements.append(
                 Measurement(
                     labelA=ref,
-                    labelB=n,
-                    DG=f_i * u,
-                    uncertainty=df_i * u,
+                    labelB=label,
+                    DG=f_i * unit,
+                    uncertainty=df_i * unit,
                     computational=True,
                     source=source,
                 )
@@ -260,8 +267,8 @@ def _estimate(
             Measurement(
                 labelA=ReferenceState(),
                 labelB=ref,
-                DG=Quantity(0.1, units=u),
-                uncertainty=Quantity(0.0, units=u),
+                DG=Quantity(0.1, units=unit),
+                uncertainty=Quantity(0.0, units=unit),
                 computational=True,
                 source=source,
             )
@@ -272,89 +279,163 @@ def _estimate(
             ligand_order=ligand_order,
         )
 
+    @staticmethod
+    def mle(
+        graph: nx.DiGraph, edge_data_label: str = "f_ij", node_data_label: str | None = None
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """
+        Compute maximum likelihood estimate of free energies and covariance in their estimates.
+        The number 'edge_data_label' is the node attribute on which the MLE will be calculated,
+        where d'edge_data_label' will be used as the standard error of the edge_data_label
 
-def _build_graph_from_measurements(
-    measurements: list[Measurement],
-) -> tuple[nx.DiGraph, object]:
-    """Build a legacy graph from the list of measurements for use in the MLE method, this is copied over from the
-    to_legacy_graph method of FEMap.
+        Reference : https://pubs.acs.org/doi/abs/10.1021/acs.jcim.9b00528
+        Xu, Huafeng. "Optimal measurement network of pairwise differences."
+        Journal of Chemical Information and Modeling 59.11 (2019): 4720-4728.
 
-    Parameters
-    ----------
-    measurements : list[Measurement]
-        Mix of relative computational and absolute experimental measurements.
-
-    Returns
-    -------
-    g : nx.DiGraph
-        Input graph ready for stats.mle
-    u : unit
-        The unit shared by all measurements (validated to be consistent).
-
-    Raises
-    ------
-    ValueError
-        If measurements have mixed units or duplicate computational edges
-        exist between the same pair of nodes.
-    """
-    if not measurements:
-        raise ValueError("No measurements provided")
-
-    units = {m.DG.u for m in measurements}
-    if len(units) > 1:
-        raise ValueError(f"All measurements must share the same units before running an estimator. Found: {units}")
-    u = next(iter(units))
-
-    g = nx.DiGraph()
-    edges_seen: list[tuple] = []
-
-    for m in measurements:
-        if not m.computational:
-            continue
-        if isinstance(m.labelA, ReferenceState):
-            continue
-        # cast to string as hashable does not support < > comparisons
-        edge_name = tuple(sorted([str(m.labelA), str(m.labelB)]))
-        if edge_name in edges_seen:
-            raise ValueError(
-                f"Multiple edges detected between nodes {m.labelA} and {m.labelB}. "
-                "MLE cannot be performed on graphs with multiple edges between the "
-                "same nodes. The results should be combined into a single estimate "
-                "and uncertainty before performing MLE. "
-                "See https://cinnabar.openfree.energy/en/latest/concepts/estimators.html"
-                "#limitations for more details."
+        NOTE: Self-edges (edges that connect a node to itself) will be ignored.
+
+        Parameters
+        ----------
+        graph :nx.Graph
+            The graph for which an estimate is to be computed
+            Each edge must have attributes 'f_ij' and 'df_ij' for the free energy and uncertainty
+            estimate
+        edge_data_label : string, default = 'f_ij'
+            edge data label of the calculated data for MLE
+        node_data_label : string, default = None
+            optional - provide if there is node data (i.e. absolute values) 'f_i' or 'exp_DG' to
+            include will expect a corresponding uncertainty 'f_di' or 'exp_dDG'
+        Returns
+        -------
+        f_i : np.array with shape (n_ligands,)
+            f_i[i] is the absolute free energy of ligand i in kcal/mol
+
+        C : np.array with shape (n_ligands, n_ligands)
+            C[i,j] is the covariance of the free energy estimates of i and j
+
+        """
+        # if we have bidirectional edge results we need to raise an error as they can not be used with MLE
+        # track the edges we have seen
+        edges = []
+        for a, b in graph.edges:
+            edge_name = (a, b) if str(a) < str(b) else (b, a)
+            edges.append(edge_name)
+
+        n_nodes = graph.number_of_nodes()
+
+        node_label = None if node_data_label is None else node_data_label.replace("_", "_d")
+        node_name_to_index = {name: i for i, name in enumerate(graph.nodes())}
+        # Adapted from the multibind implementation (Kenney IM & Beckstein O, 2023)
+        # https://github.com/Becksteinlab/multibind/blob/7c93f605d99ff67d9adef890c9302bccd2caa1b5/multibind/multibind.py#L319
+        # to support harmonic wells around individual states
+        z = np.zeros((n_nodes,))
+        F_matrix = np.zeros((n_nodes, n_nodes))
+
+        # single node harmonic wells
+        for n, data in graph.nodes(data=True):
+            if node_label in data:
+                i = node_name_to_index[n]
+                z[i] = data[node_data_label] / (data[node_label] ** 2)
+                F_matrix[i, i] = 1 / (data[node_label] ** 2)
+
+        # harmonic edge restraints
+        for a, b, data in graph.edges(data=True):
+            # self edges are ignored without warning to the user
+            if a == b:
+                continue
+
+            i = node_name_to_index[a]
+            j = node_name_to_index[b]
+
+            deltaij = data[edge_data_label]
+            if (varij := data[edge_data_label.replace("_", "_d")] ** 2) == 0:
+                raise ValueError(
+                    f"MLE solver will fail with zero reported uncertainty for calculated differences. Edge ({a}, {b}) has zero uncertainty check inputs."
+                )
+
+            z[i] += -deltaij / varij
+            z[j] += deltaij / varij
+
+            F_matrix[i, i] += 1 / varij
+            F_matrix[j, j] += 1 / varij
+            F_matrix[i, j] += -1 / varij
+            F_matrix[j, i] += -1 / varij
+
+        Finv = np.linalg.pinv(F_matrix, hermitian=True)
+        f_i = np.matmul(Finv, z)
+        return f_i, Finv
+
+    @staticmethod
+    def _build_graph_from_measurements(
+        measurements: list[Measurement],
+    ) -> tuple[nx.DiGraph, object]:
+        """Build a legacy graph from the list of measurements for use in the MLE method, this is copied over from the
+        to_legacy_graph method of FEMap.
+
+        Parameters
+        ----------
+        measurements : list[Measurement]
+            Mix of relative computational and absolute experimental measurements.
+
+        Returns
+        -------
+        graph : nx.DiGraph
+            Input graph ready for stats.mle
+        unit : unit
+            The unit shared by all measurements (validated to be consistent).
+
+        Raises
+        ------
+        ValueError
+            If measurements have mixed units or duplicate computational edges
+            exist between the same pair of nodes.
+        """
+        if not measurements:
+            raise ValueError("No measurements provided")
+
+        if len(units := {m.DG.u for m in measurements}) > 1:
+            raise ValueError(f"All measurements must share the same units before running an estimator. Found: {units}")
+        unit = units.pop()
+
+        graph = nx.DiGraph()
+        edges_seen: list[tuple] = []
+
+        # populate the edges of the graph along with their computational binding free energies
+        for m in filter(lambda m: m.computational, measurements):
+            if isinstance(m.labelA, ReferenceState):
+                # TODO this is never hit in the tests and should be supported behavior
+                continue
+            edge_name = (m.labelA, m.labelB) if str(m.labelA) < str(m.labelB) else (m.labelB, m.labelA)
+            if edge_name in edges_seen:
+                # TODO this is a limitation of the software, not the method. Support for multiple edges should be a priority
+                raise ValueError(
+                    f"Multiple edges detected between nodes {m.labelA} and {m.labelB}. "
+                    "MLE cannot be performed on graphs with multiple edges between the "
+                    "same nodes. The results should be combined into a single estimate "
+                    "and uncertainty before performing MLE. "
+                    "See https://cinnabar.openfree.energy/en/latest/concepts/estimators.html"
+                    "#limitations for more details."
+                )
+            graph.add_edge(
+                m.labelA,
+                m.labelB,
+                calc_DDG=m.DG.magnitude,
+                calc_dDDG=m.uncertainty.magnitude,
             )
-        g.add_edge(
-            m.labelA,
-            m.labelB,
-            calc_DDG=m.DG.magnitude,
-            calc_dDDG=m.uncertainty.magnitude,
-        )
-        edges_seen.append(edge_name)
-
-    # annotate nodes with experimental absolute values
-    for m in measurements:
-        if m.computational:
-            continue
-        if not isinstance(m.labelA, ReferenceState):
-            continue
-        node = m.labelB
-        if node not in g.nodes:
-            continue
-        g.nodes[node]["exp_DG"] = m.DG.magnitude
-        g.nodes[node]["exp_dDG"] = m.uncertainty.magnitude
-        g.nodes[node]["name"] = node
-
-    # infer experimental DDG for edges where both endpoints have absolute data
-    for A, B, d in g.edges(data=True):
-        try:
-            DG_A = g.nodes[A]["exp_DG"]
-            dDG_A = g.nodes[A]["exp_dDG"]
-            DG_B = g.nodes[B]["exp_DG"]
-            dDG_B = g.nodes[B]["exp_dDG"]
-        except KeyError:
-            continue
-        d["exp_DDG"] = DG_B - DG_A
-        d["exp_dDDG"] = (dDG_A**2 + dDG_B**2) ** 0.5
-
-    return g, u
+            edges_seen.append(edge_name)
+
+        # annotate nodes with experimental absolute values, this doesn't add edges
+        for m in filter(lambda m: not m.computational, measurements):
+            # labelA must always be a reference state, otherwise it is ignored
+            if not isinstance(m.labelA, ReferenceState):
+                # TODO this is never hit in the tests
+                continue
+            # TODO to support experimental values, we need this to not be true
+            # do not include experimental information if no computation data is already present
+            if (node := m.labelB) not in graph.nodes:
+                continue
+            graph.nodes[node]["exp_DG"] = m.DG.magnitude
+            graph.nodes[node]["exp_dDG"] = m.uncertainty.magnitude
+            graph.nodes[node]["name"] = node
+
+        return graph, unit

cinnabar/femap.py

@@ -21,8 +21,8 @@
 import pandas as pd
 from openff.units import Quantity, unit
 
-from cinnabar import stats
 from cinnabar.conversion import convert_observable
+from cinnabar.estimators import MLEEstimator
 from cinnabar.measurements import Measurement, ReferenceState
 
 if TYPE_CHECKING:
@@ -872,7 +872,7 @@ def to_legacy_graph(self) -> nx.DiGraph:
                 d["exp_dDDG"] = (dDG_A**2 + dDG_B**2) ** 0.5
         # apply MLE for calculated DG values
         if self.check_weakly_connected():
-            f_i_calc, C_calc = stats.mle(g, factor="calc_DDG")
+            f_i_calc, C_calc = MLEEstimator.mle(g, edge_data_label="calc_DDG")
             variance = np.diagonal(C_calc)
             variance = variance**0.5