Add functions to calculate cycle closure

cinnabar/femap.py

@@ -9,6 +9,7 @@
 
 import copy
 import itertools
+import math
 import pathlib
 import warnings
 from dataclasses import asdict
@@ -946,3 +947,155 @@ def draw_graph(
         else:
             fig.savefig(filename, bbox_inches="tight", dpi=300)
         plt.close(fig)
+
+    def get_cycle_closure_dataframe(self, max_cycle_length: int = 5) -> pd.DataFrame:
+        """
+        Calculate cycle closure errors for all cycles in the network.
+
+        Parameters
+        ----------
+        max_cycle_length : int, default 5
+            Only consider cycles up to this length. Default 5.
+
+        Returns
+        -------
+        The pandas DataFrame will have the following columns:
+        - source
+        - cycle
+        - cc (kcal/mol)
+        - cc_per_edge (kcal/mol)
+        - cc_unc_normalized
+        Sorted by source and cycle closure error descending.
+
+        Notes
+        -----
+        Three cycle closure metrics are calculated:
+
+        - ``cc (kcal/mol)``: the raw absolute sum of DDGs around the cycle. Units: kcal/mol.
+
+        - ``cc_per_edge (kcal/mol)``: the cycle closure divided by the square root of the cycle
+          length, to allow comparison across different cycle lengths;
+          see Baumann et al. (DOI 10.1021/acs.jctc.3c00282). Units: kcal/mol.
+
+        - ``cc_unc_normalized``: the cycle closure error divided by its propagated uncertainty,
+          calculated as ``abs(sum_ddgs) / sqrt(sum_var)``.
+        """
+        df = self.get_relative_dataframe()
+        comp_df = df[df["computational"]]
+
+        rows = []
+        for source, source_df in comp_df.groupby("source"):
+            edge_ddg = {(row["labelA"], row["labelB"]): row["DDG (kcal/mol)"] for _, row in source_df.iterrows()}
+            edge_uncertainty = {
+                (row["labelA"], row["labelB"]): row["uncertainty (kcal/mol)"] for _, row in source_df.iterrows()
+            }
+
+            network = nx.DiGraph()
+            for a, b in edge_ddg:
+                network.add_edge(a, b)
+
+            cycles = [c for c in nx.simple_cycles(network.to_undirected()) if len(c) <= max_cycle_length]
+
+            for cycle in cycles:
+                sum_ddgs = 0.0
+                sum_var = 0.0
+                for i, lig in enumerate(cycle):
+                    lig_a = lig
+                    lig_b = cycle[i + 1] if i < 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:
+                    cc = abs(sum_ddgs)
+                    # Normalize by sqrt(cycle length) to allow comparison across
+                    # different cycle lengths
+                    cc_per_edge = cc / math.sqrt(len(cycle))
+                    cc_z_score = cc / math.sqrt(sum_var) if sum_var > 0 else np.nan
+                    rows.append(
+                        {
+                            "source": source,
+                            "cycle": tuple(cycle),
+                            "cc (kcal/mol)": cc,
+                            "cc_per_edge (kcal/mol)": cc_per_edge,
+                            "cc_unc_normalized": cc_z_score,
+                        }
+                    )
+
+        return (
+            pd.DataFrame(
+                rows,
+                columns=["source", "cycle", "cc (kcal/mol)", "cc_per_edge (kcal/mol)", "cc_unc_normalized"],
+            )
+            .sort_values(["source", "cc (kcal/mol)"], ascending=[True, False])
+            .reset_index(drop=True)
+        )
+
+    def get_cycle_closure_edge_statistics_dataframe(self, max_cycle_length: int = 5) -> pd.DataFrame:
+        """
+        For each simulated edge, report how many cycles it appears in and
+        the mean and max cycle closure error of those cycles per source.
+
+        The cycle closure values are based on ``cc_per_edge (kcal/mol)``,
+        defined as the absolute cycle closure divided by the square root of the cycle length.
+
+        Parameters
+        ----------
+        max_cycle_length : int, default 5
+            Only consider cycles up to this length. Defaults to 5.
+
+        Returns
+        -------
+        The pandas DataFrame will have the following columns:
+        - source
+        - ligandA
+        - ligandB
+        - n_cycles
+        - mean_cc_per_edge (kcal/mol)
+        - max_cc_per_edge (kcal/mol)
+
+        Sorted by source and mean cycle closure error descending.
+        """
+        from collections import defaultdict
+
+        cc_df = self.get_cycle_closure_dataframe(max_cycle_length=max_cycle_length)
+
+        rows = []
+        for source, source_cc_df in cc_df.groupby("source"):
+            edge_cycles: dict[tuple, list[float]] = defaultdict(list)
+
+            for _, row in source_cc_df.iterrows():
+                cycle = list(row["cycle"])
+                cc_per_edge = row["cc_per_edge (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 = self._canonical_edge((lig_a, lig_b))
+                    edge_cycles[edge].append(cc_per_edge)
+
+            for (a, b), ccs in edge_cycles.items():
+                rows.append(
+                    {
+                        "source": source,
+                        "ligandA": a,
+                        "ligandB": b,
+                        "n_cycles": len(ccs),
+                        "mean_cc_per_edge (kcal/mol)": sum(ccs) / len(ccs),
+                        "max_cc_per_edge (kcal/mol)": max(ccs),
+                    }
+                )
+
+        return (
+            pd.DataFrame(rows)
+            .sort_values(["source", "mean_cc_per_edge (kcal/mol)"], ascending=[True, False])
+            .reset_index(drop=True)
+        )

cinnabar/plotting.py

@@ -1058,3 +1058,73 @@ def ecdf_plot_all_DDGs(
         filename=filename,
         **kwargs,
     )
+
+
+def plot_cycle_closure(
+    fe_map: FEMap,
+    filename: str | None,
+    max_cycle_length: int = 5,
+    sources: list[str] | None = None,
+    bin_width: float = 0.5,
+) -> plt.Figure:
+    """
+    Plot a histogram of cycle closure errors, taking the ``cc_per_edge (kcal/mol)``
+    which is the cycle closure divided by the square root of the cycle length.
+
+    Parameters
+    ----------
+    fe_map : FEMap
+        FEMap object containing the calculated edges.
+    filename : str | None, default None
+        If provided, the plot will be saved to this filename.
+    max_cycle_length : int, default 5
+        Only consider cycles up to this length. Defaults to 5.
+    sources : list[str] | None, default None
+        List of sources to plot. If None, all sources are plotted.
+    bin_width : float, default 0.5
+        Width of histogram bins in kcal/mol. Default: 0.5
+
+    Returns
+    -------
+    plt.Figure
+        The matplotlib Figure object containing the histogram which can be edited further.
+
+    Raises
+    ------
+    ValueError
+        If the FEMap contains no cycles, or if a requested
+        source cannot be found.
+    """
+    df = fe_map.get_cycle_closure_dataframe(max_cycle_length=max_cycle_length)
+
+    if df.empty:
+        raise ValueError("The FEMap does not contain cycles.")
+
+    if sources is not None:
+        df = df[df["source"].isin(sources)]
+        if df.empty:
+            raise ValueError(f"No cycles found for sources {sources}.")
+
+    unique_sources = df["source"].unique()
+
+    fig, ax = plt.subplots(figsize=(5, 4))
+
+    max_val = df["cc_per_edge (kcal/mol)"].max()
+    bins = np.arange(0, max_val + bin_width, bin_width).tolist()
+
+    for source in unique_sources:
+        source_df = df[df["source"] == source]
+        ax.hist(source_df["cc_per_edge (kcal/mol)"], bins=bins, alpha=0.6, label=source)
+
+    ax.set_xlabel(r"Cycle closure per edge (kcal mol$^{-1}$)")
+    ax.set_ylabel("Count")
+    ax.set_title("Cycle closure distribution")
+    ax.legend()
+    fig.tight_layout()
+
+    if filename is None:
+        plt.show()
+    else:
+        fig.savefig(filename, bbox_inches="tight", dpi=300)
+
+    return fig

cinnabar/tests/conftest.py

@@ -524,3 +524,25 @@ def ecdf_femap_missing_exp_data():
         label="ligand2", value=-5.0 * unit.kilocalories_per_mole, uncertainty=0.2 * unit.kilocalories_per_mole
     )
     return fe_map
+
+
+@pytest.fixture()
+def perfect_cycle():
+    """A perfect cycle with zero cycle closure."""
+    kcalpm = unit.kilocalorie_per_mole
+    fe = FEMap()
+    fe.add_relative_calculation("A", "B", value=1.0 * kcalpm, uncertainty=0.1 * kcalpm, source="method_a")
+    fe.add_relative_calculation("B", "C", value=1.0 * kcalpm, uncertainty=0.1 * kcalpm, source="method_a")
+    fe.add_relative_calculation("C", "A", value=-2.0 * kcalpm, uncertainty=0.1 * kcalpm, source="method_a")
+    return fe
+
+
+@pytest.fixture()
+def imperfect_cycle():
+    """An imperfect cycle with a known cycle closure error."""
+    kcalpm = unit.kilocalorie_per_mole
+    fe = FEMap()
+    fe.add_relative_calculation("A", "B", value=1.0 * kcalpm, uncertainty=0.1 * kcalpm, source="method_b")
+    fe.add_relative_calculation("B", "C", value=1.0 * kcalpm, uncertainty=0.1 * kcalpm, source="method_b")
+    fe.add_relative_calculation("C", "A", value=-1.5 * kcalpm, uncertainty=0.1 * kcalpm, source="method_b")
+    return fe

cinnabar/tests/test_femap.py

@@ -1,4 +1,5 @@
 import json
+import math
 import re
 
 import matplotlib.pyplot as plt
@@ -838,3 +839,60 @@ 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")
+
+
+def test_get_cycle_closure_perfect_cycle(perfect_cycle):
+    result = perfect_cycle.get_cycle_closure_dataframe()
+    assert isinstance(result, pd.DataFrame)
+    assert list(result.columns) == ["source", "cycle", "cc (kcal/mol)", "cc_per_edge (kcal/mol)", "cc_unc_normalized"]
+    assert len(result) == 1
+    assert result["cc (kcal/mol)"].iloc[0] == pytest.approx(0.0, abs=1e-6)
+    assert result["cc_per_edge (kcal/mol)"].iloc[0] == pytest.approx(0.0, abs=1e-6)
+    assert result["cc_unc_normalized"].iloc[0] == pytest.approx(0.0, abs=1e-6)
+
+
+def test_get_cycle_closure_hystereses(imperfect_cycle):
+    result = imperfect_cycle.get_cycle_closure_dataframe()
+    expected_cc = abs(0.5)
+    expected_cc_per_edge = round(abs(0.5) / math.sqrt(3), 2)
+    expected_cc_normalized = round(abs(0.5) / math.sqrt(3 * 0.1**2), 2)
+    assert result["cc (kcal/mol)"].iloc[0] == pytest.approx(expected_cc, abs=0.01)
+    assert result["cc_per_edge (kcal/mol)"].iloc[0] == pytest.approx(expected_cc_per_edge, abs=0.01)
+    assert result["cc_unc_normalized"].iloc[0] == pytest.approx(expected_cc_normalized, abs=0.01)
+
+
+def test_get_cycle_closure_multiple_sources(perfect_cycle, imperfect_cycle):
+    fe = perfect_cycle + imperfect_cycle
+
+    result = fe.get_cycle_closure_dataframe()
+    assert len(result) == 2
+    assert set(result["source"].unique()) == {"method_a", "method_b"}
+    assert result[result["source"] == "method_a"]["cc (kcal/mol)"].iloc[0] == pytest.approx(0.0, abs=1e-6)
+    assert result[result["source"] == "method_b"]["cc (kcal/mol)"].iloc[0] == pytest.approx(0.5, abs=0.01)
+
+
+def test_get_cc_based_edge_statistics_known_value(perfect_cycle):
+    result = perfect_cycle.get_cycle_closure_edge_statistics_dataframe()
+    assert len(result) == 3
+    assert (result["n_cycles"] == 1).all()
+    assert (result["mean_cc_per_edge (kcal/mol)"] == 0.0).all()
+    assert (result["max_cc_per_edge (kcal/mol)"] == 0.0).all()
+
+
+def test_get_cc_based_edge_statistics_reverse_direction(perfect_cycle):
+    kcalpm = unit.kilocalorie_per_mole
+    # add edges for more cycles
+    perfect_cycle.add_relative_calculation("A", "D", value=1.0 * kcalpm, uncertainty=0.1 * kcalpm, source="method_a")
+    perfect_cycle.add_relative_calculation("D", "B", value=-1.5 * kcalpm, uncertainty=0.1 * kcalpm, source="method_a")
+
+    result = perfect_cycle.get_cycle_closure_edge_statistics_dataframe(max_cycle_length=3)
+    cc_per_edge_abc = 0.0 / math.sqrt(3)  # perfect cycle A -> B -> C -> A
+    cc_per_edge_bad = 1.5 / math.sqrt(3)  # imperfect cycle B -> A -> D -> B
+
+    # A to B in two cycles
+    ab_row = result[(result["ligandA"] == "A") & (result["ligandB"] == "B")]
+    assert ab_row["n_cycles"].iloc[0] == 2
+    assert ab_row["mean_cc_per_edge (kcal/mol)"].iloc[0] == pytest.approx(
+        (cc_per_edge_abc + cc_per_edge_bad) / 2, abs=1e-3
+    )
+    assert ab_row["max_cc_per_edge (kcal/mol)"].iloc[0] == pytest.approx(cc_per_edge_bad, abs=1e-3)

cinnabar/tests/test_plotting.py

@@ -511,3 +511,52 @@ def test_plot_ecdf_colors(fe_map, tmp_path):
     # check that the line color matches the specified color
     line = fig.get_axes()[0].lines[0]
     assert line.get_color() == "#FF5733"
+
+
+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 per edge (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_dataframe().empty
+    output_file = tmp_path / "cycle_closure.png"
+    with pytest.raises(ValueError, match="The FEMap does not contain cycles"):
+        plotting.plot_cycle_closure(fe, filename=str(output_file))
+    assert not output_file.exists()
+
+
+def test_plot_cycle_closure_invalid_source(fe_map, tmp_path):
+    output_file = tmp_path / "cycle_closure.png"
+    with pytest.raises(ValueError, match="No cycles found for sources"):
+        plotting.plot_cycle_closure(fe_map, filename=str(output_file), sources=["nonexistent_source"])
+    assert not output_file.exists()
+
+
+def test_plot_cycle_closure_multiple_sources(perfect_cycle, imperfect_cycle, tmp_path):
+    fe = perfect_cycle + imperfect_cycle
+    output_file = tmp_path / "cycle_closure_multiple.png"
+    fig = plotting.plot_cycle_closure(fe, filename=str(output_file))
+    assert fig is not None
+    axes = fig.get_axes()[0]
+    legend_texts = [t.get_text() for t in axes.get_legend().get_texts()]
+    assert "method_a" in legend_texts
+    assert "method_b" in legend_texts
+    assert output_file.exists()