Add functions to calculate cycle closure

cinnabar/femap.py

@@ -8,6 +8,7 @@
 """
 
 import copy
+import math
 import itertools
 import pathlib
 import warnings
@@ -745,3 +746,117 @@ def draw_graph(self, title: str = "", filename: Union[str, None] = None):
             plt.show()
         else:
             plt.savefig(filename, bbox_inches="tight")
+
+    def get_cycle_closure(self, max_cycle_length: int = 5) -> list[tuple[str, float]]:
+        """
+        Calculate cycle closure errors for all cycles in the network.
+
+        Parameters
+        ----------
+        max_cycle_length : int, optional
+            Only consider cycles up to this length. Default 5.
+
+        Returns
+        -------
+        pd.DataFrame
+            DataFrame with columns 'cycle', 'cc (kcal/mol)', sorted by
+            cycle closure error descending.
+        """
+        network = self.to_legacy_graph()
+        edge_ddg = {(a, b): d["calc_DDG"] for a, b, d in network.edges(data=True)}
+        edge_uncertainty = {(a, b): d["calc_dDDG"] for a, b, d in network.edges(data=True)}
+
+        # Find all ligand cycles
+        cycles = [c for c in nx.simple_cycles(network.to_undirected()) if len(c) <= max_cycle_length]
+
+        # Loop over cycles, calculate sum of DG along cycle
+        rows = []
+        for cycle in cycles:
+            # Store DDG values along the cycle
+            sum_ddgs = 0.0
+            sum_var = 0.0
+            for inx, ligand in enumerate(cycle):
+                lig_a = ligand
+                lig_b = cycle[inx + 1] if inx < len(cycle) - 1 else cycle[0]
+
+                # depending on the direction the edge was calculated,
+                # the sign of the DDG has to change
+                if (lig_a, lig_b) in edge_ddg:
+                    sum_ddgs += edge_ddg[(lig_a, lig_b)]
+                    sum_var += edge_uncertainty[(lig_a, lig_b)] ** 2
+                elif (lig_b, lig_a) in edge_ddg:
+                    sum_ddgs -= edge_ddg[(lig_b, lig_a)]
+                    sum_var += edge_uncertainty[(lig_b, lig_a)] ** 2
+                else:
+                    # Edge missing from network; skip this cycle
+                    break
+
+            else:
+                # Normalize by sqrt(cycle length) to allow comparison across
+                # different cycle lengths
+                cc = abs(sum_ddgs / math.sqrt(len(cycle)))
+                cc_uncertainty_normalized = abs(sum_ddgs) / math.sqrt(sum_var)
+                rows.append(
+                    {
+                        "cycle": tuple(cycle),
+                        "cc (kcal/mol)": round(cc, 2),
+                        "cc_unc_normalized (kcal/mol)": round(cc_uncertainty_normalized, 2),
+                    }
+                )
+
+        df = (
+            pd.DataFrame(rows, columns=["cycle", "cc (kcal/mol)", "cc_unc_normalized (kcal/mol)"])
+            .sort_values("cc (kcal/mol)", ascending=False)
+            .reset_index(drop=True)
+        )
+
+        return df
+
+    def get_cc_based_edge_statistics(self, max_cycle_length: int = 5) -> pd.DataFrame:
+        """
+        For each edge, report how many cycles it appears in and
+        the mean and max cycle closure error of those cycles.
+        Edges with high mean closure error across many cycles are
+        likely candidates for re-simulation.
+
+        Parameters
+        ----------
+        max_cycle_length : int, optional
+            Only consider cycles up to this length. Defaults to 5.
+
+        Returns
+        -------
+        pd.DataFrame
+            DataFrame with columns 'ligandA', 'ligandB', 'n_cycles',
+            'mean_cc (kcal/mol)', 'max_cc (kcal/mol)', sorted by
+            mean cycle closure error descending.
+        """
+        from collections import defaultdict
+
+        cc_df = self.get_cycle_closure(max_cycle_length=max_cycle_length)
+        network = self.to_legacy_graph()
+        edge_ddg = {(a, b): d["calc_DDG"] for a, b, d in network.edges(data=True)}
+
+        edge_cycles: dict[tuple, list[float]] = defaultdict(list)
+        for _, row in cc_df.iterrows():
+            cycle = list(row["cycle"])
+            cc = row["cc (kcal/mol)"]
+            for i, lig in enumerate(cycle):
+                lig_a = lig
+                lig_b = cycle[i + 1] if i < len(cycle) - 1 else cycle[0]
+                edge = (lig_a, lig_b) if (lig_a, lig_b) in edge_ddg else (lig_b, lig_a)
+                edge_cycles[edge].append(cc)
+
+        rows = []
+        for (a, b), ccs in edge_cycles.items():
+            rows.append(
+                {
+                    "ligandA": a,
+                    "ligandB": b,
+                    "n_cycles": len(ccs),
+                    "mean_cc (kcal/mol)": round(sum(ccs) / len(ccs), 3),
+                    "max_cc (kcal/mol)": round(max(ccs), 3),
+                }
+            )
+
+        return pd.DataFrame(rows).sort_values("mean_cc (kcal/mol)", ascending=False).reset_index(drop=True)

cinnabar/plotting.py

@@ -1,4 +1,5 @@
 import itertools
+import warnings
 from typing import Any, Literal, Optional, Union
 
 import matplotlib.pylab as plt
@@ -893,3 +894,45 @@ def ecdf_plot_all_DDGs(
         **kwargs,
     )
     return fig
+
+
+def plot_cycle_closure(
+    fe_map: FEMap,
+    filename: Optional[str] = None,
+    max_cycle_length: int = 5,
+) -> plt.Figure:
+    """
+    Plot a histogram of cycle closure errors.
+
+    Parameters
+    ----------
+    fe_map : FEMap
+        FEMap object containing the calculated edges.
+    filename : str, optional
+        If provided, the plot will be saved to this filename.
+    max_cycle_length : int, optional
+        Only consider cycles up to this length. Defaults to 5.
+        The matplotlib Figure object, which can be edited further.
+    """
+    df = fe_map.get_cycle_closure(max_cycle_length=max_cycle_length)
+
+    if df.empty:
+        warnings.warn("No cycles found; skipping plot.")
+        return
+
+    errors = df["cc (kcal/mol)"]
+    n_cycles = len(errors)
+
+    fig, ax = plt.subplots(figsize=(5, 4))
+    ax.hist(errors, bins="auto", alpha=0.6, color="steelblue")
+    ax.set_xlabel(r"Cycle closure (kcal mol$^{-1}$)")
+    ax.set_ylabel("Count")
+    ax.set_title(f"Cycle closure distribution (n={n_cycles})")
+    fig.tight_layout()
+
+    if filename is None:
+        plt.show()
+    else:
+        fig.savefig(filename, bbox_inches="tight", dpi=300)
+
+    return fig

cinnabar/tests/test_femap.py

@@ -1,4 +1,5 @@
 import json
+import math
 
 import matplotlib.pyplot as plt
 import networkx as nx
@@ -659,3 +660,44 @@ def test_missing_estimator_metadata(example_map):
     with pytest.raises(KeyError, match="No estimator metadata stored for source test."):
         example_map.generate_absolute_values()
         example_map.get_estimator_metadata("test")
+
+
+@pytest.fixture()
+def perfect_cycle():
+    """A perfect cycle with zero cycle closure."""
+    kcalpm = unit.kilocalorie_per_mole
+    fe = femap.FEMap()
+    fe.add_relative_calculation("A", "B", value=1.0 * kcalpm, uncertainty=0.1 * kcalpm)
+    fe.add_relative_calculation("B", "C", value=1.0 * kcalpm, uncertainty=0.1 * kcalpm)
+    fe.add_relative_calculation("C", "A", value=-2.0 * kcalpm, uncertainty=0.1 * kcalpm)
+    return fe
+
+
+def test_get_cycle_closure_known_value(perfect_cycle):
+    result = perfect_cycle.get_cycle_closure()
+    assert isinstance(result, pd.DataFrame)
+    assert list(result.columns) == ["cycle", "cc (kcal/mol)", "cc_unc_normalized (kcal/mol)"]
+    assert len(result) == 1
+    assert result["cc (kcal/mol)"].iloc[0] == pytest.approx(0.0, abs=1e-6)
+    assert result["cc_unc_normalized (kcal/mol)"].iloc[0] == pytest.approx(0.0, abs=1e-6)
+
+
+def test_get_cycle_closure_normalized_known_value():
+    kcalpm = unit.kilocalorie_per_mole
+    fe = femap.FEMap()
+    fe.add_relative_calculation("A", "B", value=1.0 * kcalpm, uncertainty=0.1 * kcalpm)
+    fe.add_relative_calculation("B", "C", value=1.0 * kcalpm, uncertainty=0.1 * kcalpm)
+    fe.add_relative_calculation("C", "A", value=-1.5 * kcalpm, uncertainty=0.1 * kcalpm)
+
+    result = fe.get_cycle_closure()
+    expected_cc_normalized = round(abs(0.5) / math.sqrt(3 * 0.1**2), 2)
+
+    assert result["cc_unc_normalized (kcal/mol)"].iloc[0] == pytest.approx(expected_cc_normalized, abs=0.01)
+
+
+def test_get_cc_based_edge_statistics_known_value(perfect_cycle):
+    result = perfect_cycle.get_cc_based_edge_statistics()
+    assert len(result) == 3
+    assert (result["n_cycles"] == 1).all()
+    assert (result["mean_cc (kcal/mol)"] == 0.0).all()
+    assert (result["max_cc (kcal/mol)"] == 0.0).all()

cinnabar/tests/test_plotting.py

@@ -417,3 +417,34 @@ def test_plot_ecdf_dgs_no_data(graph):
         match="Graph with label test has nodes with missing calculated DG values, which should be stored as `calc_DG`.",
     ):
         plotting.ecdf_plot_DGs([graph], labels=["test"], filename=None)
+
+
+def test_plot_cycle_closure(fe_map, tmp_path):
+    output_file = tmp_path / "cycle_closure.png"
+    fig = plotting.plot_cycle_closure(fe_map, filename=str(output_file))
+    assert fig is not None
+    axes = fig.get_axes()[0]
+    assert axes.get_xlabel() == r"Cycle closure (kcal mol$^{-1}$)"
+    assert axes.get_ylabel() == "Count"
+    assert output_file.exists()
+
+
+def test_plot_cycle_closure_show(fe_map, show_called):
+    _ = plotting.plot_cycle_closure(fe_map, filename=None)
+    assert "show" in show_called
+
+
+def test_plot_cycle_closure_no_cycles_no_plot(tmp_path):
+    fe = FEMap()
+    fe.add_relative_calculation(
+        "A",
+        "B",
+        value=1.0 * unit.kilocalorie_per_mole,
+        uncertainty=0.1 * unit.kilocalorie_per_mole,
+    )
+    assert fe.get_cycle_closure().empty
+    output_file = tmp_path / "cycle_closure.png"
+    with pytest.warns(UserWarning, match="No cycles found"):
+        fig = plotting.plot_cycle_closure(fe, filename=str(output_file))
+    assert fig is None
+    assert not output_file.exists()